def load_optimizer(self):
if self.arg.optimizer == 'SGD':
self.optimizer = optim.SGD(self.model.parameters(),
lr=self.arg.base_lr,
momentum=0.9,
nesterov=self.arg.nesterov,
weight_decay=self.arg.weight_decay)
optimor = optim.SGD
elif self.arg.optimizer == 'Adam':
self.optimizer = optim.Adam(self.model.parameters(),
lr=self.arg.base_lr,
weight_decay=self.arg.weight_decay)
elif self.arg.optimizer == 'Adamod':
self.optimizer = adamod.AdaMod(self.model.parameters(),
lr=self.arg.base_lr,
beta3=0.999)
print("I am using Adamod")
else:
raise ValueError()
model = STModel(encoder_layer_size=param['encoder_layer_size'],
decoder_layer_size=param['decoder_layer_size'],
kernel_size=param['kernel_size'],
out_channels=param['filter_size'],
in_channels=train_data.x.shape[1],
input_width=train_data.x.shape[3],
input_height=train_data.x.shape[4],
hidden_size=hidden_size,
prediction_window=prediction_window,
device=device).to(device)
criterion = WeightedRMSELoss()
#optimizer_params = {'lr': 0.001}
#optimizer = torch.optim.Adam(net.parameters(), **optimizer_params)
opt_params = {'lr': 0.001, 'beta3': 0.999}
optimizer = adamod.AdaMod(model.parameters(), **opt_params)
model
# In[17]:
model_path = os.path.join('../../../models/CHIRPS/Reconstructions/ST-RFD' +
'_' + datetime.now().strftime('m%md%d-h%Hm%Ms%S') +
'.pth.tar')
trainer = Trainer(model,
train_loader,
val_loader,
criterion,
optimizer,
100,
device,
model_path,