Build entire model before predicting

Author: Colin Grambow

reacdiff/parsing.py

@@ -15,16 +15,53 @@ def parse_predict_args():
     parser = argparse.ArgumentParser(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     )
+
+    # General arguments
     parser.add_argument('--data_path', type=str, required=True,
                         help='Path to data containing states for prediction task')
     parser.add_argument('--model', type=str, required=True,
                         help='Path to trained model')
+    parser.add_argument('--output_dim', type=int, required=True,
+                        help='Dimensionality of predictions')
     parser.add_argument('--data_path2', type=str,
                         help='Path to additional observable states for prediction')
     parser.add_argument('--save_path', type=str, default=os.path.join(os.getcwd(), 'preds.csv'),
                         help='Path to save predictions to')
+    parser.add_argument('--cpu', action='store_true',
+                        help='Run on CPU instead of GPU')
     parser.add_argument('--batch_size', type=int, default=32,
                         help='Batch size')
+
+    # Encoder arguments
+    parser.add_argument('--feat_maps', type=int, default=16,
+                        help='Number of feature maps in first convolutional layer')
+    parser.add_argument('--first_conv_size', type=int, default=7,
+                        help='Filter size in first convolutional layer')
+    parser.add_argument('--first_conv_stride', type=int, default=2,
+                        help='Strides in first convolutional layer')
+    parser.add_argument('--first_pool_size', type=int, default=3,
+                        help='Window size in first pool')
+    parser.add_argument('--first_pool_stride', type=int, default=2,
+                        help='Strides in first pool')
+    parser.add_argument('--growth_rate', type=int, default=12,
+                        help='Growth rate in dense blocks')
+    parser.add_argument('--blocks', type=int, nargs='+', default=[3, 4, 5],
+                        help='Numbers of layers in each dense block')
+    parser.add_argument('--dropout', type=float, default=0.0,
+                        help='Dropout rate after convolutions')
+    parser.add_argument('--reduction', type=float, default=0.5,
+                        help='Compression rate in transition blocks')
+    parser.add_argument('--no_bottleneck', action='store_true',
+                        help='Do not use bottleneck convolution in dense blocks')
+    parser.add_argument('--flatten_last', action='store_true',
+                        help='Flatten instead of global pool before output')
+
+    # RNN arguments
+    parser.add_argument('--rnn_layers', type=int, default=1,
+                        help='Number of RNN layers')
+    parser.add_argument('--rnn_units', type=int, default=100,
+                        help='Number of units in RNN')
+
     return parser.parse_args()
 
 

reacdiff/train/predict.py

@@ -1,9 +1,7 @@
 import os
 
-import keras
-
 import reacdiff.data.data as datamod
-import reacdiff.utils as utils
+import reacdiff.models.crnn as models
 
 
 def predict(args):
@@ -16,9 +14,32 @@ def predict(args):
 
     os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
 
-    # Load model
-    model = keras.models.load_model(args.model, custom_objects={'rmse': utils.rmse, 'mae': utils.mae})
+    # Load model (currently have to build entire model with provided arguments)
+    # Build model
+    crnn = models.CRNN(
+        time_steps=data.data.shape[1],
+        input_shape=data.data.shape[2:],
+        output_dim=args.output_dim,
+        observables=data.get_num_observables(),
+        rnn_layers=args.rnn_layers,
+        rnn_units=args.rnn_units,
+        use_gpu=not args.cpu
+    )
+    crnn.build(
+        feat_maps=args.feat_maps,
+        first_conv_size=args.first_conv_size,
+        first_conv_stride=args.first_conv_stride,
+        first_pool_size=args.first_pool_size,
+        first_pool_stride=args.first_pool_stride,
+        growth_rate=args.growth_rate,
+        blocks=args.blocks,
+        dropout=args.dropout,
+        reduction=args.reduction,
+        bottleneck=not args.no_bottleneck,
+        flatten_last=args.flatten_last
+    )
+    crnn.model.load_weights(args.model)
 
     # Predict
-    preds = model.predict(data.get_data(), batch_size=args.batch_size, verbose=1)
+    preds = crnn.model.predict(data.get_data(), batch_size=args.batch_size, verbose=1)
     datamod.save_csv(preds, args.save_path)