fix: all in 1 deps install, amp option fix #32

Author: ydcjeff

requirements-dev.txt

@@ -1,3 +1,5 @@
+-r requirements.txt
+
 # dev
 pytorch-ignite
 torch

templates/_base/_argparse.pyi

@@ -4,6 +4,10 @@ from argparse import ArgumentParser
 
 {% block defaults %}
 DEFAULTS = {
+    "use_amp": {
+        "action": "store_true",
+        "help": "use torch.cuda.amp for automatic mixed precision"
+    },
     "seed": {
         "default": 666,
         "type": int,

templates/_base/_handlers.pyi

@@ -157,6 +157,7 @@ def get_logger(
         Config object for setting up loggers
 
     `config` has to contain
+    - `filepath`: logging path to output file
     - `logger_log_every_iters`: logging iteration interval for loggers
 
     train_engine
@@ -232,5 +233,6 @@ def get_logger(
         **kwargs,
     )
     {% else %}
-    return None
+    logger_handler = None
     {% endif %}
+    return logger_handler

templates/single/_sidebar.py

@@ -19,7 +19,7 @@ def get_configs():
     default_none_options(config)
 
     with st.beta_expander("Single Model, Single Optimizer Template Configurations", expanded=True):
-        st.info("Those in the parenthesis are used in the generated code.")
+        st.info("Names in the parenthesis are variable names used in the generated code.")
 
         # group by configurations type
         distributed_options(config)

templates/single/single_cg/engines.pyi

@@ -2,22 +2,24 @@
 `train_engine` and `eval_engine` like trainer and evaluator
 """
 from typing import Any, Tuple
+
 import torch
 from ignite.engine import Engine
+from torch.cuda.amp import autocast
 from torch.optim.optimizer import Optimizer
 
 from single_cg.events import TrainEvents, train_events_to_attr
 
 
 # Edit below functions the way how the model will be training
 
-# train_fn is how the model will be learning with given batch
-# below in the train_fn, common parameters are provided
+# train_function is how the model will be learning with given batch
+# below in the train_function, common parameters are provided
 # you can add any additional parameters depending on the training
 # NOTE : engine and batch parameters are needed to work with
 # Ignite's Engine.
 # TODO: Extend with your custom training.
-def train_fn(
+def train_function(
     config: Any,
     engine: Engine,
     batch: Any,
@@ -26,7 +28,7 @@ def train_fn(
     optimizer: Optimizer,
     device: torch.device,
 ):
-    """Training.
+    """Model training step.
 
     Parameters
     ----------
@@ -48,8 +50,9 @@ def train_fn(
     samples = batch[0].to(device, non_blocking=True)
     targets = batch[1].to(device, non_blocking=True)
 
-    outputs = model(samples)
-    loss = loss_fn(outputs, targets)
+    with autocast(enabled=config.use_amp):
+        outputs = model(samples)
+        loss = loss_fn(outputs, targets)
 
     loss.backward()
     engine.state.backward_completed += 1
@@ -66,22 +69,22 @@ def train_fn(
     return loss_value
 
 
-# evaluate_fn is how the model will be learning with given batch
-# below in the evaluate_fn, common parameters are provided
+# evaluate_function is how the model will be learning with given batch
+# below in the evaluate_function, common parameters are provided
 # you can add any additional parameters depending on the training
 # NOTE : engine and batch parameters are needed to work with
 # Ignite's Engine.
 # TODO: Extend with your custom evaluation.
 @torch.no_grad()
-def evaluate_fn(
+def evaluate_function(
     config: Any,
     engine: Engine,
     batch: Any,
     model: torch.nn.Module,
     loss_fn: torch.nn.Module,
     device: torch.device,
 ):
-    """Evaluating.
+    """Model evaluating step.
 
     Parameters
     ----------
@@ -102,10 +105,11 @@ def evaluate_fn(
     samples = batch[0].to(device, non_blocking=True)
     targets = batch[1].to(device, non_blocking=True)
 
-    outputs = model(samples)
-    loss = loss_fn(outputs, targets)
-    loss_value = loss.item()
+    with autocast(enabled=config.use_amp):
+        outputs = model(samples)
+        loss = loss_fn(outputs, targets)
 
+    loss_value = loss.item()
     engine.state.metrics = {"eval_loss": loss_value}
     return loss_value
 
@@ -117,21 +121,21 @@ def create_engines(**kwargs) -> Tuple[Engine, Engine]:
 
     Parameters
     ----------
-    kwargs: keyword arguments passed to both train_fn and evaluate_fn
+    kwargs: keyword arguments passed to both train_function and evaluate_function
 
     Returns
     -------
     train_engine, eval_engine
     """
     train_engine = Engine(
-        lambda e, b: train_fn(
+        lambda e, b: train_function(
             engine=e,
             batch=b,
             **kwargs,
         )
     )
     eval_engine = Engine(
-        lambda e, b: evaluate_fn(
+        lambda e, b: evaluate_function(
             engine=e,
             batch=b,
             **kwargs,

templates/single/tests/test_engines.py

@@ -1,11 +1,12 @@
 import unittest
+from argparse import Namespace
 from numbers import Number
 from unittest.mock import MagicMock
 
 import ignite.distributed as idist
 import torch
 from ignite.engine.engine import Engine
-from single_cg.engines import create_engines, evaluate_fn, train_fn
+from single_cg.engines import create_engines, evaluate_function, train_function
 from single_cg.events import TrainEvents, train_events_to_attr
 from torch import nn, optim
 
@@ -27,7 +28,8 @@ def test_train_fn(self):
         optim = MagicMock()
         engine.add_event_handler(TrainEvents.BACKWARD_COMPLETED, backward)
         engine.add_event_handler(TrainEvents.OPTIM_STEP_COMPLETED, optim)
-        output = train_fn(None, engine, self.batch, self.model, self.loss_fn, self.optimizer, self.device)
+        config = Namespace(use_amp=False)
+        output = train_function(config, engine, self.batch, self.model, self.loss_fn, self.optimizer, self.device)
         self.assertIsInstance(output, Number)
         self.assertTrue(hasattr(engine.state, "backward_completed"))
         self.assertTrue(hasattr(engine.state, "optim_step_completed"))
@@ -39,7 +41,10 @@ def test_train_fn(self):
         self.assertTrue(optim.called)
 
     def test_train_fn_event_filter(self):
-        engine = Engine(lambda e, b: train_fn(None, e, b, self.model, self.loss_fn, self.optimizer, self.device))
+        config = Namespace(use_amp=False)
+        engine = Engine(
+            lambda e, b: train_function(config, e, b, self.model, self.loss_fn, self.optimizer, self.device)
+        )
         engine.register_events(*TrainEvents, event_to_attr=train_events_to_attr)
         backward = MagicMock()
         optim = MagicMock()
@@ -60,7 +65,10 @@ def test_train_fn_event_filter(self):
         self.assertTrue(optim.called)
 
     def test_train_fn_every(self):
-        engine = Engine(lambda e, b: train_fn(None, e, b, self.model, self.loss_fn, self.optimizer, self.device))
+        config = Namespace(use_amp=False)
+        engine = Engine(
+            lambda e, b: train_function(config, e, b, self.model, self.loss_fn, self.optimizer, self.device)
+        )
         engine.register_events(*TrainEvents, event_to_attr=train_events_to_attr)
         backward = MagicMock()
         optim = MagicMock()
@@ -77,7 +85,10 @@ def test_train_fn_every(self):
         self.assertTrue(optim.called)
 
     def test_train_fn_once(self):
-        engine = Engine(lambda e, b: train_fn(None, e, b, self.model, self.loss_fn, self.optimizer, self.device))
+        config = Namespace(use_amp=False)
+        engine = Engine(
+            lambda e, b: train_function(config, e, b, self.model, self.loss_fn, self.optimizer, self.device)
+        )
         engine.register_events(*TrainEvents, event_to_attr=train_events_to_attr)
         backward = MagicMock()
         optim = MagicMock()
@@ -95,12 +106,13 @@ def test_train_fn_once(self):
 
     def test_evaluate_fn(self):
         engine = Engine(lambda e, b: 1)
-        output = evaluate_fn(None, engine, self.batch, self.model, self.loss_fn, self.device)
+        config = Namespace(use_amp=False)
+        output = evaluate_function(config, engine, self.batch, self.model, self.loss_fn, self.device)
         self.assertIsInstance(output, Number)
 
     def test_create_engines(self):
         train_engine, eval_engine = create_engines(
-            config=None,
+            config=Namespace(use_amp=True),
             model=self.model,
             loss_fn=self.loss_fn,
             optimizer=self.optimizer,

templates/single/tests/test_handlers.py

@@ -52,30 +52,31 @@ def test_get_handlers(self):
             self.assertIsInstance(timer_handler, (type(None), Timer), "Shoulde be Timer or None")
 
     def test_get_logger(self):
-        config = Namespace(logger_log_every_iters=1)
-        train_engine = Engine(lambda e, b: b)
-        optimizer = optim.Adam(nn.Linear(1, 1).parameters())
-        logger_handler = get_logger(
-            config=config,
-            train_engine=train_engine,
-            eval_engine=train_engine,
-            optimizers=optimizer,
-        )
-        self.assertIsInstance(
-            logger_handler,
-            (
-                BaseLogger,
-                ClearMLLogger,
-                MLflowLogger,
-                NeptuneLogger,
-                PolyaxonLogger,
-                TensorboardLogger,
-                VisdomLogger,
-                WandBLogger,
-                type(None),
-            ),
-            "Should be Ignite provided loggers or None",
-        )
+        with TemporaryDirectory() as tmp:
+            config = Namespace(filepath=tmp, logger_log_every_iters=1)
+            train_engine = Engine(lambda e, b: b)
+            optimizer = optim.Adam(nn.Linear(1, 1).parameters())
+            logger_handler = get_logger(
+                config=config,
+                train_engine=train_engine,
+                eval_engine=train_engine,
+                optimizers=optimizer,
+            )
+            self.assertIsInstance(
+                logger_handler,
+                (
+                    BaseLogger,
+                    ClearMLLogger,
+                    MLflowLogger,
+                    NeptuneLogger,
+                    PolyaxonLogger,
+                    TensorboardLogger,
+                    VisdomLogger,
+                    WandBLogger,
+                    type(None),
+                ),
+                "Should be Ignite provided loggers or None",
+            )
 
 
 if __name__ == "__main__":