Add teacher forcing ratio

Author: djosix

config.py

@@ -20,5 +20,6 @@
     'window_size': 200,
     'stride_size': 10,
     'use_transposition': False,
-    'control_ratio': 1.0
+    'control_ratio': 1.0,
+    'teacher_forcing_ratio': 1.0
 }

model.py

@@ -98,7 +98,7 @@ def expand_controls(self, controls, steps):
         return controls.repeat(steps, 1, 1)
 
     def generate(self, init, steps, events=None, controls=None, greedy=1.0,
-                 temperature=1.0, output_type='index', verbose=False):
+                 temperature=1.0, teacher_forcing_ratio=1.0, output_type='index', verbose=False):
         # init [batch_size, init_dim]
         # events [steps, batch_size] indeces
         # controls [1 or steps, batch_size, control_dim]
@@ -142,7 +142,8 @@ def generate(self, init, steps, events=None, controls=None, greedy=1.0,
                 assert False
 
             if use_teacher_forcing and step < steps - 1: # avoid last one
-                event = events[step].unsqueeze(0)
+                if np.random.random() <= teacher_forcing_ratio:
+                    event = events[step].unsqueeze(0)
 
         return torch.cat(outputs, 0)
 

train.py

@@ -68,6 +68,11 @@ def get_options():
                       type='float',
                       default=config.train['control_ratio'])
 
+    parser.add_option('-T', '--teacher-forcing-ratio',
+                      dest='teacher_forcing_ratio',
+                      type='float',
+                      default=config.train['teacher_forcing_ratio'])
+
     parser.add_option('-t', '--use-transposition',
                       dest='use_transposition',
                       action='store_true',
@@ -99,6 +104,7 @@ def get_options():
 stride_size = options.stride_size
 use_transposition = options.use_transposition
 control_ratio = options.control_ratio
+teacher_forcing_ratio = options.teacher_forcing_ratio
 reset_optimizer = options.reset_optimizer
 
 event_dim = EventSeq.dim()
@@ -121,6 +127,7 @@ def get_options():
 print('Window size:', window_size)
 print('Stride size:', stride_size)
 print('Control ratio:', control_ratio)
+print('Teacher forcing ratio:', teacher_forcing_ratio)
 print('Random transposition:', use_transposition)
 print('Reset optimizer:', reset_optimizer)
 print('Device:', device)
@@ -201,7 +208,8 @@ def save_model():
             controls = None
 
         init = torch.randn(batch_size, model.init_dim).to(device)
-        outputs = model.generate(init, window_size, events[:-1], controls, output_type='logit')
+        outputs = model.generate(init, window_size, events=events[:-1], controls=controls,
+                                 teacher_forcing_ratio=teacher_forcing_ratio, output_type='logit')
         assert outputs.shape[:2] == events.shape[:2]
 
         loss = loss_function(outputs.view(-1, event_dim), events.view(-1))