I used these:
###training loop, scheduler learning was used to find the best learning rate. in evert 250 epoch learning rate changes by the factor of 0.1
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=320, gamma=0.7, last_epoch=-1)
epoch_num = 0
train_error = []
test_error = []
best_error = 100
for epoch in range(num_epochs):
print(scheduler.get_last_lr())
loss_total = 0
test_loss_total = 0
model.train()
for batch_idx, sample in enumerate(train_loader):
inp = sample
inp = inp.cuda()
output = model(inp).cuda()
loss = criterion(output, inp)
loss_total += loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
loss_total = loss_total / (batch_idx+1)
train_error.append(loss_total.item())
model.eval()
with torch.no_grad():
for batch_idx, sample in enumerate(test_loader):
inp = sample
inp = inp.cuda()
output = model(inp).cuda()
loss = criterion(output, inp)
test_loss_total += loss
test_loss_total = test_loss_total / (batch_idx+1)
if loss < best_error:
best_error = loss
best_epoch = epoch
print(‘Best loss at epoch’, best_epoch)
model_save_name = ‘2020-05-25 Li_sample_no_image_processing (5)’
path = F"/content/drive/My Drive/{model_save_name}"
torch.save(model.state_dict(), path)
test_error.append(test_loss_total.item())
if epoch%10 == 9:
epoch_num+=epoch_num
print ((’\r Train Epoch : {}/{} \tLoss : {:.4f}’.format (epoch+1,num_epochs,loss_total)))
print ((’\r Test Epoch : {}/{} \tLoss : {:.4f}’.format (epoch+1,num_epochs,test_loss_total)))
I do print the avg value over mini batch size
Thank you Patrick