update Lightning module test for OSS

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/17

Use PyTorch Lightning checkpoint in the test.

Reviewed By: zhanghang1989

Differential Revision: D26962697

fbshipit-source-id: abe635e374c3ada130243f0eaadff34204f04fa1

tests/test_lightning_train_net.py

@@ -41,10 +41,10 @@ class TestLightningTrainNet(unittest.TestCase):
             ckpts,
         )
 
-        with tempfile.TemporaryDirectory() as tmp_dir2:
+        tmp_dir2 = tempfile.TemporaryDirectory()  # noqa to avoid flaky test
         cfg2 = cfg.clone()
         cfg2.defrost()
-            cfg2.OUTPUT_DIR = tmp_dir2
+        cfg2.OUTPUT_DIR = tmp_dir2.name
         # load the last checkpoint from previous training
         cfg2.MODEL.WEIGHTS = os.path.join(tmp_dir, "last.ckpt")
 
@@ -53,3 +53,4 @@ class TestLightningTrainNet(unittest.TestCase):
         accuracy2 = flatten_config_dict(out2.accuracy)
         for k in accuracy:
             np.testing.assert_equal(accuracy[k], accuracy2[k])
+        tmp_dir2.cleanup()

tests/test_runner_lightning_task.py

@@ -8,22 +8,19 @@ import unittest
 from copy import deepcopy
 from typing import Dict
 
-import d2go.runner.default_runner as default_runner
 import pytorch_lightning as pl  # type: ignore
 import torch
 from d2go.config import CfgNode
 from d2go.runner.lightning_task import GeneralizedRCNNTask
+from d2go.tests import meta_arch_helper as mah
 from detectron2.utils.events import EventStorage
+from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
 from torch import Tensor
 
-from d2go.tests import meta_arch_helper as mah
-
-OSSRUN = os.getenv('OSSRUN') == '1'
 
 class TestLightningTask(unittest.TestCase):
     def _get_cfg(self, tmp_dir: str) -> CfgNode:
-        runner = default_runner.Detectron2GoRunner()
-        cfg = mah.create_detection_cfg(runner, tmp_dir)
+        cfg = mah.create_detection_cfg(GeneralizedRCNNTask, tmp_dir)
         cfg.TEST.EVAL_PERIOD = cfg.SOLVER.MAX_ITER
         return cfg
 
@@ -38,14 +35,11 @@ class TestLightningTask(unittest.TestCase):
                 return False
         return True
 
-    @unittest.skipIf(OSSRUN, "not supported yet")
     def test_load_from_checkpoint(self) -> None:
         with tempfile.TemporaryDirectory() as tmp_dir:
             task = GeneralizedRCNNTask(self._get_cfg(tmp_dir))
-            from stl.lightning.callbacks.model_checkpoint import ModelCheckpoint
-            checkpoint_callback = ModelCheckpoint(
-                directory=task.cfg.OUTPUT_DIR, has_user_data=False
-            )
+
+            checkpoint_callback = ModelCheckpoint(dirpath=task.cfg.OUTPUT_DIR)
             params = {
                 "max_steps": 1,
                 "limit_train_batches": 1,
@@ -92,15 +86,13 @@ class TestLightningTask(unittest.TestCase):
                 self._compare_state_dict(init_state, task.ema_state.state_dict())
             )
 
-    @unittest.skipIf(OSSRUN, "not supported yet")
     def test_load_ema_weights(self):
         with tempfile.TemporaryDirectory() as tmp_dir:
             cfg = self._get_cfg(tmp_dir)
             cfg.MODEL_EMA.ENABLED = True
             task = GeneralizedRCNNTask(cfg)
-            from stl.lightning.callbacks.model_checkpoint import ModelCheckpoint
             checkpoint_callback = ModelCheckpoint(
-                directory=task.cfg.OUTPUT_DIR, save_last=True
+                dirpath=task.cfg.OUTPUT_DIR, save_last=True
             )
 
             trainer = pl.Trainer(
@@ -115,9 +107,19 @@ class TestLightningTask(unittest.TestCase):
                 trainer.fit(task)
 
             # load EMA weights from checkpoint
-            task2 = GeneralizedRCNNTask.load_from_checkpoint(os.path.join(tmp_dir, "last.ckpt"))
-            self.assertTrue(self._compare_state_dict(task.ema_state.state_dict(), task2.ema_state.state_dict()))
+            task2 = GeneralizedRCNNTask.load_from_checkpoint(
+                os.path.join(tmp_dir, "last.ckpt")
+            )
+            self.assertTrue(
+                self._compare_state_dict(
+                    task.ema_state.state_dict(), task2.ema_state.state_dict()
+                )
+            )
 
             # apply EMA weights to model
             task2.ema_state.apply_to(task2.model)
-            self.assertTrue(self._compare_state_dict(task.ema_state.state_dict(), task2.model.state_dict()))
+            self.assertTrue(
+                self._compare_state_dict(
+                    task.ema_state.state_dict(), task2.model.state_dict()
+                )
+            )