Modify predictions to work with ensembles

Author: Colin Grambow

reacdiff/parsing.py

@@ -19,6 +19,10 @@ def parse_predict_args():
                         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('--targets_path', type=str,
+                        help='Path to targets')
+    parser.add_argument('--targets_size', type=int, default=4,
+                        help='Size of target vector. Only required if targets_path is not specified.')
     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'),

reacdiff/train/predict.py

@@ -1,6 +1,9 @@
+import glob
 import os
+import re
 
 import keras
+import numpy as np
 
 import reacdiff.data.data as datamod
 import reacdiff.utils as utils
@@ -11,14 +14,40 @@ def predict(args):
     print('Loading data')
     data = datamod.Dataset(
         datamod.load_data(args.data_path),
+        targets=None if args.targets_path is None else datamod.load_csv(args.targets_path),
         data2=None if args.data_path2 is None else datamod.load_data(args.data_path2)
     )
 
-    os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
+    os.makedirs(os.path.dirname(os.path.abspath(args.save_path)), exist_ok=True)
 
-    # Load model
-    model = keras.models.load_model(args.model, custom_objects={'rmse': utils.rmse, 'mae': utils.mae})
+    # Walk directory for ensemble of models
+    if os.path.isdir(args.model):
+        model_dirs = glob.iglob(os.path.join(args.model, 'model*'))
+        model_nums = [re.search('\d+', os.path.basename(d))[0] for d in model_dirs]
+        model_nums.sort()
 
-    # Predict
-    preds = model.predict(data.get_data(), batch_size=args.batch_size, verbose=1)
+        targets_size = args.targets_size if args.targets_path is None else data.targets.shape[1]
+        all_preds = np.zeros((len(model_nums), len(data), targets_size))
+
+        for model_idx in model_nums:
+            print(f'Evaluating model {model_idx}')
+
+            model_path = os.path.join(args.model, f'model{model_idx}', 'model.h5')
+            model = keras.models.load_model(model_path, custom_objects={'rmse': utils.rmse, 'mae': utils.mae})
+
+            preds = model.predict(data.get_data(), batch_size=args.batch_size, verbose=1)
+            all_preds[model_idx] = preds
+        preds = np.mean(all_preds, axis=0)
+    else:
+        # Load model
+        model = keras.models.load_model(args.model, custom_objects={'rmse': utils.rmse, 'mae': utils.mae})
+
+        # Predict
+        preds = model.predict(data.get_data(), batch_size=args.batch_size, verbose=1)
+
+    if args.targets_path is not None:
+        print('Evaluating ensemble')
+        rmse = utils.rmse_np(data.targets, preds)
+        mae = utils.mae_np(data.targets, preds)
+        print(f'rmse: {rmse:.4f}; mae: {mae:.4f}')
     datamod.save_csv(preds, args.save_path)