Merge pull request amathislab#5 from amathislab/new-spindle

New spindle model and task (for revised version)

Author: AlexEMG

code/kinematics_decoding.py

@@ -1,6 +1,7 @@
 import os
 import sys
 from nn_models import ConvModel, RecurrentModel, AffineModel
+from nn_rmodels import ConvRModel, RecurrentRModel
 
 import numpy as np
 import pandas as pd
@@ -39,6 +40,74 @@ def load_model(meta_data, experiment_id, model_type, is_trained):
             t_stride=int(meta_data['t_stride']),
             seed=myseed,
             train=is_trained)
+    
+    if model_type == 'rec':
+        
+        try:
+            myseed = meta_data['seed']
+        except:
+            myseed = None
+        
+        model = RecurrentModel(
+            experiment_id=experiment_id,
+            nclasses=20,
+            rec_blocktype=meta_data['rec_blocktype'],
+            n_recunits=meta_data['n_recunits'],
+            npplayers=meta_data['npplayers'],
+            s_kernelsize=meta_data['s_kernelsize'],
+            s_stride=meta_data['s_stride'],
+            nppfilters=meta_data['nppfilters'],
+            seed=myseed,
+            train=is_trained)
+
+    return model
+
+
+def load_rmodel(meta_data, experiment_id, model_type, is_trained):
+    '''Load a trained `ConvRModel`, `RecurrentRModel` object.
+
+    Returns
+    -------
+    model : an instance of `ConvRModel` or `RecurrentRModel`
+
+    '''
+    if model_type == 'conv':
+        
+        try:
+            myseed = meta_data['seed']
+        except:
+            myseed = None
+        
+        model = ConvRModel(
+            experiment_id=experiment_id,
+            arch_type=meta_data['arch_type'],
+            nlayers=meta_data['nlayers'],
+            n_skernels=meta_data['n_skernels'],
+            n_tkernels=meta_data['n_tkernels'],
+            s_kernelsize=int(meta_data['s_kernelsize']),
+            t_kernelsize=int(meta_data['t_kernelsize']),
+            s_stride=int(meta_data['s_stride']),
+            t_stride=int(meta_data['t_stride']),
+            seed=myseed,
+            train=is_trained)
+    
+    if model_type == 'rec':
+        
+        try:
+            myseed = meta_data['seed']
+        except:
+            myseed = None
+        
+        model = RecurrentRModel(
+            experiment_id=experiment_id,
+            rec_blocktype=meta_data['rec_blocktype'],
+            n_recunits=meta_data['n_recunits'],
+            npplayers=meta_data['npplayers'],
+            s_kernelsize=meta_data['s_kernelsize'],
+            s_stride=meta_data['s_stride'],
+            nppfilters=meta_data['nppfilters'],
+            seed=myseed,
+            train=is_trained)
 
     return model
 
@@ -81,7 +150,7 @@ def set_kin_dimensions(kinematics, n_out_timesteps):
     assert n_timesteps == 320, "Kinematics shape mismatch. Please revise."
     skip_idx = n_timesteps // (n_out_timesteps - 1)
     kinematics = kinematics[:, :, ::skip_idx]
-    assert kinematics.shape[-1] == n_out_timesteps, "Whaatt!"
+    assert kinematics.shape[-1] == n_out_timesteps
     return kinematics.transpose(0, 2, 1)
 
 
@@ -141,7 +210,11 @@ def make_h5file_name(model_path, repfile_name, layername):
 
     with tf.Graph().as_default():
         X = tf.placeholder(tf.float32, myshape, name="X")
-        _, _, output = model.predict(X, is_training=False)
+        name = type(model).__name__
+        if name == 'ConvModel' or name == 'RecurrentModel':
+            _, _, output = model.predict(X, is_training=False)
+        elif name == 'ConvRModel' or name == 'RecurrentRModel':
+            _, output = model.predict(X, is_training=False)
         restorer = tf.train.Saver()
         myconfig = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
 

code/nn_models.py

@@ -244,7 +244,8 @@ class RecurrentModel():
     """Defines a recurrent neural network model of the proprioceptive system."""
 
     def __init__(
-            self, experiment_id, nclasses, rec_blocktype, n_recunits, npplayers, nppunits, keep_prob):
+            self, experiment_id, nclasses, rec_blocktype, n_recunits, npplayers, nppfilters, 
+            s_kernelsize, s_stride, seed=None, train=True, CPU=False):
         """Set up the hyperparameters of the recurrent model.
 
         Arguments
@@ -254,29 +255,37 @@ def __init__(
         rec_blocktype: {'lstm', 'gru'} str, type of recurrent block.
         n_recunits : int, number of units in the recurrent block.
         npplayers : int, number of layers in the fully-connected module.
-        nppunits : list of ints, number of units in the affine layers for spatial processing.
-        keep_prob : float, amount of dropout at each spatial processing layer.
+        nppfilters : list of ints, number of filters (spatial convolutions) for spatial processing.
+        s_kernelsize : int, size of conv kernel
+        s_stride : int, stride for conv kernel
+        seed : int, for saving random initializations
+        train : bool, whether to train the model or not (just save random initialization)
 
         """
 
-        assert len(nppunits) == npplayers
-        assert rec_blocktype in ('lstm', 'gru')
+        assert len(nppfilters) == npplayers
+        assert rec_blocktype in  ('lstm', 'gru')
 
         self.experiment_id = experiment_id
         self.nclasses = nclasses
         self.rec_blocktype = rec_blocktype
         self.n_recunits = n_recunits
         self.npplayers = npplayers
-        self.nppunits = nppunits
-        self.keep_prob = keep_prob
+        self.nppfilters = nppfilters
+        self.s_kernelsize = s_kernelsize
+        self.s_stride = s_stride
+        self.seed = seed
+        self.CPU = CPU
 
         # Make model name
-        dropout_name = {'0.6': 'high', '0.7': 'med', '0.8': 'low'}
-        units = ('-'.join(str(i) for i in nppunits))
-        parts_name = [rec_blocktype, str(npplayers), units, str(n_recunits), dropout_name[str(keep_prob)]]
+        units = ('-'.join(str(i) for i in nppfilters))
+        parts_name = [rec_blocktype, str(npplayers), units, str(n_recunits)]
 
         # Create model directory
         self.name = '_'.join(parts_name)
+        if seed is not None: self.name += '_' + str(self.seed)
+        if not train: self.name += 'r'
+
         exp_dir = os.path.join(MODELS_DIR, f'experiment_{self.experiment_id}')
         self.model_path = os.path.join(exp_dir, self.name)
 
@@ -292,25 +301,33 @@ def predict(self, X, is_training=True):
         net = OrderedDict()
 
         with tf.variable_scope('Network', reuse=tf.AUTO_REUSE):
-            score = tf.transpose(X, [0, 2, 1])
+            score = X
             batch_size = X.get_shape()[0]
 
-            fully_connected = lambda score, layer_id: slim.fully_connected(
-                score, self.nppunits[layer_id], normalizer_fn=slim.layer_norm, scope=f'FC{layer_id}')
+            spatial_conv = lambda score, layer_id: slim.conv2d(
+                score, self.nppfilters[layer_id], [self.s_kernelsize, 1], [self.s_stride, 1],
+                data_format='NHWC', normalizer_fn=slim.layer_norm, scope=f'Spatial{layer_id}')
 
             for layer in range(self.npplayers):
-                score = fully_connected(score, layer)
-                score = slim.dropout(score, keep_prob=self.keep_prob, is_training=is_training)
+                score = spatial_conv(score, layer)
                 net[f'spatial{layer}'] = score
 
             # `cudnn_rnn` requires the inputs to be of shape [timesteps, batch_size, num_inputs]
+            score = tf.transpose(score, [0, 2, 1, 3])
+            score = tf.reshape(score, [batch_size, 320, -1])
             score = tf.transpose(score, [1, 0, 2])
-            if self.rec_blocktype == 'lstm':
+            
+            if self.rec_blocktype == 'lstm' and self.CPU == False:
                 recurrent_cell = cudnn_rnn.CudnnLSTM(1, self.n_recunits, name='RecurrentBlock')
                 score, _ = recurrent_cell.apply(score)
-            elif self.rec_blocktype == 'gru':
+            elif self.rec_blocktype == 'gru' and self.CPU == False:
                 recurrent_cell = cudnn_rnn.CudnnGRU(1, self.n_recunits, name='RecurrentBlock')
                 score, _ = recurrent_cell.apply(score)
+            elif self.rec_blocktype == 'lstm' and self.CPU:
+                with tf.variable_scope('RecurrentBlock'):
+                    rec_layer = lambda: cudnn_rnn.CudnnCompatibleLSTMCell(self.n_recunits)
+                recurrent_cell = tf.nn.rnn_cell.MultiRNNCell([rec_layer() for _ in range(1)])
+                score, _ = tf.nn.dynamic_rnn(recurrent_cell, score, dtype=tf.float32)
 
             score = tf.transpose(score, [1, 0, 2])
             net['recurrent_out'] = score