Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

added option to change default tensor

This commit is contained in:
William Falcon 2019-05-13 21:52:02 -04:00
parent 90a460ec62
commit 1c7d477d03
2 changed files with 10 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
pytorch_lightning/models/trainer.py
View File

@ -7,9 +7,6 @@ from pytorch_lightning.root_module.model_saving import TrainerIO
from torch.optim.lr_scheduler import MultiStepLR from torch.optim.lr_scheduler import MultiStepLR
import pdb import pdb
try: try:
from apex import amp from apex import amp
APEX_AVAILABLE = True APEX_AVAILABLE = True
@ -86,7 +83,6 @@ class Trainer(TrainerIO):
use_amp = True use_amp = True
self.use_amp = use_amp and APEX_AVAILABLE self.use_amp = use_amp and APEX_AVAILABLE
if self.use_amp: if self.use_amp:
self.amp_handle = amp.init(enabled=True)
print('using 16bit precision') print('using 16bit precision')
def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct): def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
@ -223,14 +219,14 @@ class Trainer(TrainerIO):
# filter out the weights that were done on gpu so we can load on good old cpus # filter out the weights that were done on gpu so we can load on good old cpus
self.optimizers = model.configure_optimizers() self.optimizers = model.configure_optimizers()
# if self.use_amp: if self.use_amp:
# An example # An example
# self.model, optimizer = amp.initialize( self.model, optimizer = amp.initialize(
# self.model, self.optimizers[0], opt_level="O2", self.model, self.optimizers[0], opt_level="O2",
# keep_batchnorm_fp32=True, loss_scale="dynamic" keep_batchnorm_fp32=True, loss_scale="dynamic"
# ) )
# self.optimizers[0] = optimizer self.optimizers[0] = optimizer
# model.trainer = self model.trainer = self
# add lr schedulers # add lr schedulers
if self.lr_scheduler_milestones is not None: if self.lr_scheduler_milestones is not None:
@ -374,7 +370,7 @@ class Trainer(TrainerIO):
# backward pass # backward pass
if self.use_amp: if self.use_amp:
for optimizer in self.optimizers: for optimizer in self.optimizers:
with self.amp_handle.scale_loss(loss, optimizer) as scaled_loss: with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward() scaled_loss.backward()
else: else:
loss.backward() loss.backward()

2
setup.py
View File

@ -7,7 +7,7 @@ from setuptools import setup, find_packages
# http://blog.ionelmc.ro/2014/05/25/python-packaging/ # http://blog.ionelmc.ro/2014/05/25/python-packaging/
setup( setup(
name="pytorch-lightning", name="pytorch-lightning",
version='0.1.dev1722', version='0.1.dev1723',
description="The Keras for ML researchers using PyTorch", description="The Keras for ML researchers using PyTorch",
author="William Falcon", author="William Falcon",
author_email="waf2107@columbia.edu", author_email="waf2107@columbia.edu",