condition and supervise

broadinstitute · Jan 15, 2025 · b287451 · b287451
1 parent 6d8d1d4
commit b287451
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/ml4h/models/train.py b/ml4h/models/train.py
@@ -379,6 +379,9 @@ def train_diffusion_control_model(args, supervised=False):
     model.load_weights(checkpoint_path)
 
     if args.inspect_model:
+        metrics = model.evaluate(generate_test, batch_size=args.batch_size, steps=args.test_steps, return_dict=True)
+        logging.info(f'Test metrics: {metrics}')
+
         data, labels, paths = big_batch_from_minibatch_generator(generate_test, 1)
         preds = model.plot_reconstructions((data, labels), prefix=f'{args.output_folder}/{args.id}/reconstructions/')
         image_out = {args.tensor_maps_in[0].output_name(): data[args.tensor_maps_in[0].input_name()]}