modified: dsm_torch.py

Author: chiragnagpal

dsm/dsm_torch.py

@@ -392,11 +392,10 @@ class DeepConvolutionalSurvivalMachinesTorch(DeepSurvivalMachinesTorch):
       Dimensionality of the input features. A tuple (height, width).
   k: int
       The number of underlying parametric distributions.
+  embedding: torch.nn.Module
+      A torch CNN to obtain the representation of the input data.
   hidden: int
       The number of neurons in each hidden layer.
-  init: tuple
-      A tuple for initialization of the parameters for the underlying
-      distributions. (shape, scale).
   dist: str
       Choice of the underlying survival distributions.
       One of 'Weibull', 'LogNormal'.
@@ -411,8 +410,8 @@ class DeepConvolutionalSurvivalMachinesTorch(DeepSurvivalMachinesTorch):
 
   """
 
-  def __init__(self, inputdim, k, typ='ConvNet',
-               hidden=None, dist='Weibull',
+  def __init__(self, inputdim, k,
+               embedding=None, hidden=None, dist='Weibull',
                temp=1000., discount=1.0, optimizer='Adam', risks=1):
     super(DeepSurvivalMachinesTorch, self).__init__()
 
@@ -427,9 +426,13 @@ def __init__(self, inputdim, k, typ='ConvNet',
 
     self._init_dsm_layers(hidden)
 
-    self.embedding = create_conv_representation(inputdim=inputdim,
-                                                hidden=hidden,
-                                                typ='ConvNet')
+    if embedding is None:
+      self.embedding = create_conv_representation(inputdim=inputdim,
+                                                  hidden=hidden,
+                                                  typ='ConvNet')
+    else:
+      self.embedding = embedding
+
 
   def forward(self, x, risk='1'):
     """The forward function that is called when data is passed through DSM.