returnn 1.20230512.225636__tar.gz → 1.20230515.100902__tar.gz

{returnn-1.20230512.225636/returnn.egg-info → returnn-1.20230515.100902}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230512.225636
+Version: 1.20230515.100902
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20230515.100902/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230515.100902'
+long_version = '1.20230515.100902+git.57b171e'

{returnn-1.20230512.225636 → returnn-1.20230515.100902}/returnn/torch/engine.py

@@ -4,10 +4,13 @@ Main engine for PyTorch
 
 from __future__ import annotations
 from typing import Optional, Union, Callable, Dict
+from contextlib import nullcontext
 
 import os
 import torch
 import torch.utils.data.datapipes as dp
+from torch import autocast
+from torch.cuda import amp
 from torchdata.dataloader2 import DataLoader2
 from random import random
 
@@ -53,9 +56,25 @@ class Engine(EngineBase):
         self._save_model_epoch_interval = 1
         self._updater = None  # type: Optional[Updater]
 
+        self._use_autocast = False
+        self._autocast_dtype = None  # type: Optional[str]
+        self._grad_scaler = None  # type: Optional[amp.GradScaler]
+
         self._device = _get_device_from_config(config)
         print("Using device:", self._device, file=log.v2)
 
+        amp_options = self.config.typed_value("torch_amp_options")
+        if amp_options:
+            assert isinstance(amp_options, dict)
+            amp_options = util.CollectionReadCheckCovered(amp_options)
+            self._use_autocast = True
+            self._autocast_dtype = amp_options["dtype"]
+            amp_options.assert_all_read()
+
+        use_grad_scaler = self.config.bool("use_grad_scaler", self._use_autocast)
+        if use_grad_scaler:
+            self._grad_scaler = amp.GradScaler()
+
     def init_train_from_config(
         self,
         config: Optional[Config] = None,
@@ -144,6 +163,7 @@ class Engine(EngineBase):
         accumulated_inv_norm_factors_dict = NumbersDict()
         step_idx = 0
         for data in self._train_dataloader:
+            self._updater.get_optimizer().zero_grad()
             self._run_step(data, train_flag=True)
 
             train_ctx = rf.get_run_ctx()
@@ -158,9 +178,13 @@ class Engine(EngineBase):
                 {name: float(_to_raw(loss.get_inv_norm_factor())) for name, loss in train_ctx.losses.items()}
             )
 
-            self._updater.get_optimizer().zero_grad()
-            total_loss.raw_tensor.backward()
-            self._updater.get_optimizer().step()
+            if self._grad_scaler is not None:
+                self._grad_scaler.scale(total_loss).backward()
+                self._grad_scaler.step(self._updater.get_optimizer())
+                self._grad_scaler.update()
+            else:
+                total_loss.raw_tensor.backward()
+                self._updater.get_optimizer().step()
 
             accumulated_losses_dict += losses_dict
             accumulated_inv_norm_factors_dict += inv_norm_factors_dict
@@ -320,8 +344,9 @@ class Engine(EngineBase):
 
         rf.init_train_step_run_ctx(train_flag=train_flag)
 
-        sentinel_kw = {"__fwd_compatible_random_arg_%i" % int(random() * 100): None}
-        self._train_step_func(model=self._orig_model, extern_data=extern_data, **sentinel_kw)
+        with autocast(device_type=self._device, dtype=self._autocast_dtype) if self._use_autocast else nullcontext():
+            sentinel_kw = {"__fwd_compatible_random_arg_%i" % int(random() * 100): None}
+            self._train_step_func(model=self._orig_model, extern_data=extern_data, **sentinel_kw)
 
     def _load_model(self, *, epoch: int):
         """

{returnn-1.20230512.225636 → returnn-1.20230515.100902/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230512.225636
+Version: 1.20230515.100902
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20230512.225636/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230512.225636'
-long_version = '1.20230512.225636+git.88127d1'