added load on CPU first (#221)

* added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added print logs * added print logs * changed close order * changed close order
2026-09-10 12:21:57 +08:00 · 2019-09-11 07:52:36 -04:00
parent 90353ac54e
commit 9576dd28b2
4 changed files with 207 additions and 51 deletions
@@ -214,10 +214,20 @@ def get_hparams(continue_training=False, hpc_exp_number=0):


 def main():
-    """Verify test() on fitted model"""
+    """
+    Make sure DDP + AMP continue training correctly
+    :return:
+    """
    hparams = get_hparams()
    model = LightningTestModel(hparams)

+    trainer_options = dict(
+        show_progress_bar=True,
+        max_nb_epochs=4,
+        gpus=2,
+        distributed_backend='dp',
+    )
+
    save_dir = init_save_dir()

    # exp file to get meta
@@ -228,31 +238,59 @@ def main():
    # exp file to get weights
    checkpoint = ModelCheckpoint(save_dir)

-    trainer_options = dict(
-        show_progress_bar=False,
-        max_nb_epochs=1,
-        train_percent_check=0.4,
-        val_percent_check=0.2,
-        checkpoint_callback=checkpoint,
-        experiment=exp,
-        gpus=[0, 1],
-        distributed_backend='ddp'
-    )
+    # add these to the trainer options
+    trainer_options['experiment'] = exp
+    trainer_options['checkpoint_callback'] = checkpoint

    # fit model
    trainer = Trainer(**trainer_options)
+    trainer.is_slurm_managing_tasks = True
    result = trainer.fit(model)

+    # track epoch before saving
+    real_global_epoch = trainer.current_epoch
+
    # correct result and ok accuracy
-    assert result == 1, 'training failed to complete'
-    pretrained_model = load_model(exp, save_dir, on_gpu=True, module_class=LightningTestModel)
+    assert result == 1, 'amp + dp model failed to complete'

+    # ---------------------------
+    # HPC LOAD/SAVE
+    # ---------------------------
+    # save
+    trainer.hpc_save(save_dir, exp)
+
+    # init new trainer
+    new_exp = get_exp(False, version=exp.version)
+    trainer_options['experiment'] = new_exp
+    trainer_options['checkpoint_callback'] = ModelCheckpoint(save_dir)
+    trainer_options['train_percent_check'] = 0.2
+    trainer_options['val_percent_check'] = 0.2
+    trainer_options['max_nb_epochs'] = 1
    new_trainer = Trainer(**trainer_options)
-    new_trainer.test(pretrained_model)

-    # test we have good test accuracy
-    assert_ok_test_acc(new_trainer)
-    # clear_save_dir()
+    # set the epoch start hook so we can predict before the model does the full training
+    def assert_good_acc():
+        assert trainer.current_epoch == real_global_epoch and trainer.current_epoch > 0
+
+        # if model and state loaded correctly, predictions will be good even though we
+        # haven't trained with the new loaded model
+        dp_model = new_trainer.model
+        dp_model.eval()
+
+        _ = [run_prediction(dataloader, dp_model, dp=True) for dataloader in trainer.val_dataloader]
+
+    # new model
+    model = LightningTestModel(hparams)
+    model.on_sanity_check_start = assert_good_acc
+
+    # fit new model which should load hpc weights
+    new_trainer.fit(model)
+
+    # test freeze on gpu
+    model.freeze()
+    model.unfreeze()
+
+    clear_save_dir()


 if __name__ == '__main__':