release v0.1.dev182

.gitignore

@@ -7,6 +7,7 @@ test_tube_data/
 datasets/
 model_weights/
 app/models/
+pip-wheel-metadata/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

pytorch_lightning/models/trainer.py

@@ -5,7 +5,13 @@ from pytorch_lightning.root_module.memory import get_gpu_memory_map
 import traceback
 from pytorch_lightning.root_module.model_saving import TrainerIO
 from torch.optim.lr_scheduler import MultiStepLR
+import pdb
 
+try:
+    from apex import amp
+    APEX_AVAILABLE = True
+except ModuleNotFoundError:
+    APEX_AVAILABLE = False
 
 class Trainer(TrainerIO):
 
@@ -26,6 +32,7 @@ class Trainer(TrainerIO):
                  train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0, val_check_interval=0.95,
                  log_save_interval=1, add_log_row_interval=1,
                  lr_scheduler_milestones=None,
+                 use_amp=False,
                  nb_sanity_val_steps=5):
 
         # Transfer params
@@ -73,6 +80,11 @@ class Trainer(TrainerIO):
         self.__determine_data_use_amount(train_percent_check, val_percent_check, test_percent_check, overfit_pct)
         print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu))
 
+        # apex test
+        self.use_amp = use_amp and APEX_AVAILABLE
+        if self.use_amp:
+            print('using 16bit precision')
+
     def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
         """
         Use less data for debugging purposes
@@ -86,7 +98,7 @@ class Trainer(TrainerIO):
             self.test_percent_check = overfit_pct
 
     def __is_function_implemented(self, f_name):
-        f_op = getattr(self, f_name, None)
+        f_op = getattr(self.model, f_name, None)
         return callable(f_op)
 
     @property
@@ -195,6 +207,7 @@ class Trainer(TrainerIO):
     # -----------------------------
     def fit(self, model):
         self.model = model
+        model.trainer = self
 
         # transfer data loaders from model
         self.__get_dataloaders(model)
@@ -206,6 +219,14 @@ class Trainer(TrainerIO):
         # filter out the weights that were done on gpu so we can load on good old cpus
         self.optimizers = model.configure_optimizers()
 
+        if self.use_amp:
+            # An example
+            self.model, optimizer = amp.initialize(
+                self.model, self.optimizers[0], opt_level="O2",
+            )
+            self.optimizers[0] = optimizer
+            model.trainer = self
+
         # add lr schedulers
         if self.lr_scheduler_milestones is not None:
             for optimizer in self.optimizers:
@@ -266,28 +287,25 @@ class Trainer(TrainerIO):
                 if met_batch_limit:
                     break
 
-                # give model a chance to end epoch early
-                if self.model.should_stop_epoch(data_batch):
-                    break
-
                 # ---------------
                 # RUN TRAIN STEP
                 # ---------------
-                self.__run_tng_batch(data_batch)
+                batch_result = self.__run_tng_batch(data_batch)
+                early_stop_epoch = batch_result == -1
 
                 # ---------------
                 # RUN VAL STEP
                 # ---------------
                 is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
-                if self.fast_dev_run or is_val_check_batch:
+                if self.fast_dev_run or is_val_check_batch or early_stop_epoch:
                     self.__run_validation()
 
                 # when batch should be saved
-                if (batch_nb + 1) % self.log_save_interval == 0:
+                if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
                     self.experiment.save()
 
                 # when metrics should be logged
-                if batch_nb % self.add_log_row_interval == 0:
+                if batch_nb % self.add_log_row_interval == 0 or early_stop_epoch:
                     # count items in memory
                     # nb_params, nb_tensors = count_mem_items()
 
@@ -311,6 +329,10 @@ class Trainer(TrainerIO):
                 if self.__is_function_implemented('on_batch_end'):
                     self.model.on_batch_end()
 
+                # end epoch early
+                if early_stop_epoch:
+                    break
+
             # hook
             if self.__is_function_implemented('on_epoch_end'):
                 self.model.on_epoch_end()
@@ -325,13 +347,16 @@ class Trainer(TrainerIO):
                 if stop:
                     return
 
+
     def __run_tng_batch(self, data_batch):
         if data_batch is None:
-            return
+            return 0
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            self.model.on_batch_start()
+            response = self.model.on_batch_start(data_batch)
+            if response == -1:
+                return -1
 
         if self.enable_tqdm:
             self.prog_bar.update(1)
@@ -342,7 +367,13 @@ class Trainer(TrainerIO):
         self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
         # backward pass
+        if self.use_amp:
+            for optimizer in self.optimizers:
+                with amp.scale_loss(loss, optimizer) as scaled_loss:
+                    scaled_loss.backward()
+        else:
             loss.backward()
+
         self.batch_loss_value += loss.item()
 
         # gradient update with accumulated gradients
@@ -373,6 +404,8 @@ class Trainer(TrainerIO):
         if self.__is_function_implemented('on_batch_end'):
             self.model.on_batch_end()
 
+        return 0
+
     def __run_validation(self):
         # decide if can check epochs
         can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0

pytorch_lightning/root_module/hooks.py

@@ -1,7 +1,7 @@
 import torch
 
 class ModelHooks(torch.nn.Module):
-    def on_batch_start(self):
+    def on_batch_start(self, data_batch):
         pass
 
     def on_batch_end(self):
@@ -19,5 +19,3 @@ class ModelHooks(torch.nn.Module):
     def on_post_performance_check(self):
         pass
 
-    def should_stop_epoch(self, data_batch):
-        return False

pytorch_lightning/root_module/model_saving.py

@@ -88,6 +88,7 @@ class TrainerIO(object):
         self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
         self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
         self.global_step = checkpoint['global_step']
+        self.current_epoch = checkpoint['epoch']
 
         # restore the optimizers
         optimizer_states = checkpoint['optimizer_states']

pytorch_lightning/root_module/root_module.py

@@ -24,6 +24,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         self.overfit = hparams.overfit
         self.gradient_clip = hparams.gradient_clip
         self.num = 2
+        self.trainer = None
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False
@@ -39,8 +40,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
 
         if self.on_gpu:
             print('running on gpu...')
-            self.dtype = torch.cuda.FloatTensor
-            torch.set_default_tensor_type('torch.cuda.FloatTensor')
+            torch.set_default_tensor_type(hparams.default_tensor_type)
 
     def forward(self, *args, **kwargs):
         """
@@ -150,18 +150,22 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
             return 0, 0
 
     @classmethod
-    def load_from_metrics(cls, weights_path, tags_csv, on_gpu):
+    def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
         """
         Primary way of loading model from csv weights path
         :param weights_path:
         :param tags_csv:
         :param on_gpu:
+        :param map_location: dic for mapping storage {'cuda:1':'cuda:0'}
         :return:
         """
         hparams = load_hparams_from_tags_csv(tags_csv)
         hparams.__setattr__('on_gpu', on_gpu)
 
         if on_gpu:
+            if map_location is not None:
+                checkpoint = torch.load(weights_path, map_location=map_location)
+            else:
                 checkpoint = torch.load(weights_path)
         else:
             checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)

pytorch_lightning/trainer_main.py

@@ -9,7 +9,6 @@ from pytorch_lightning.utils.arg_parse import add_default_args
 from time import sleep
 
 from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
-
 SEED = 2334
 torch.manual_seed(SEED)
 np.random.seed(SEED)

pytorch_lightning/utils/arg_parse.py

@@ -49,6 +49,8 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     parser.add_argument('--gpus', default='0', type=str)
     parser.add_argument('--single_run_gpu', dest='single_run_gpu', action='store_true')
     parser.add_argument('--disable_cuda', dest='disable_cuda', action='store_true')
+    parser.add_argument('--default_tensor_type', default='torch.cuda.FloatTensor', type=str)
+    parser.add_argument('--use_amp', dest='use_amp', action='store_true')
 
     # run on hpc
     parser.add_argument('--on_cluster', dest='on_cluster', action='store_true')

setup.py

@@ -1,16 +1,13 @@
 #!/usr/bin/env python
 
-from setuptools import setup, find_packages, os
+from setuptools import setup, find_packages
 
 # https://packaging.python.org/guides/single-sourcing-package-version/
-version = {}
-with open(os.path.join("src", "pytorch-lightning", "__init__.py")) as fp:
-    exec(fp.read(), version)
 
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 setup(
     name="pytorch-lightning",
-    version=version["__version__"],
+    version='0.1.dev182',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",
@@ -24,35 +21,9 @@ setup(
         "tqdm",
         "test-tube",
     ],
-    extras_require={
-        "dev": [
-            "black ; python_version>='3.6'",
-            "coverage",
-            "isort",
-            "pytest",
-            "pytest-cov<2.6.0",
-            "pycodestyle",
-            "sphinx",
-            "nbsphinx",
-            "ipython>=5.0",
-            "jupyter-client",
-        ]
-    },
-    packages=find_packages("src"),
-    package_dir={"": "src"},
-    classifiers=[
-        "Development Status :: 4 - Beta",
-        "Intended Audience :: Education",
-        "Intended Audience :: Science/Research",
-        "License :: OSI Approved :: MIT License",
-        "Operating System :: OS Independent",
-        "Programming Language :: Python",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.5",
-        "Programming Language :: Python :: 3.6",
-        "Programming Language :: Python :: 3.7",
-    ],
+    packages=find_packages(),
     long_description=open("README.md", encoding="utf-8").read(),
+    long_description_content_type='text/markdown',
     include_package_data=True,
     zip_safe=False,
 )

src/pytorch-lightning/__init__.py

@@ -1,10 +0,0 @@
-"""
-=================
-pytorch-lightning
-=================
-
-The Keras for ML researchers using PyTorch. More control. Less boilerplate.
-
-"""
-
-__version__ = "0.1.dev01"