wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-10 12:21:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add src, docs and other important folders

Browse Source
This commit is contained in:
Shreyas Bapat
2019-04-03 22:16:02 +05:30
parent f26488bd16
commit 18b0c5a122
36 changed files with 260 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/root_module/optimization.py
-22
View File
@@ -1,22 +0,0 @@
from torch import nn
from torch import optim
class OptimizerConfig(nn.Module):
def choose_optimizer(self, optimizer, params, optimizer_params, opt_name_key):
if optimizer == 'adam':
optimizer = optim.Adam(params, **optimizer_params)
if optimizer == 'sparse_adam':
optimizer = optim.SparseAdam(params, **optimizer_params)
if optimizer == 'sgd':
optimizer = optim.SGD(params, **optimizer_params)
if optimizer == 'adadelta':
optimizer = optim.Adadelta(params, **optimizer_params)
# transfer opt state if loaded
if opt_name_key in self.loaded_optimizer_states_dict:
state = self.loaded_optimizer_states_dict[opt_name_key]
optimizer.load_state_dict(state)
return optimizer