Can someone help me with this please,
def train_net(n_epochs):
valid_loss_min = np.Inf
history = {'train_loss': [], 'valid_loss': [], 'epoch': []}
for epoch in range(n_epochs):
train_loss = 0.0
valid_loss = 0.0
net.train()
running_loss = 0.0
for batch_i, data in enumerate(train_loader):
images = data['image']
key_pts = data['keypoints']
key_pts = key_pts.view(key_pts.size(0), -1)
key_pts = key_pts.type(torch.FloatTensor).to(device)
images = images.type(torch.FloatTensor).to(device)
output_pts = net(images)
loss = criterion(output_pts, key_pts)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss += loss.item()*images.data.size(0)
net.eval()
with torch.no_grad():
for batch_i, data in enumerate(test_loader):
images = data['image']
key_pts = data['keypoints']
key_pts = key_pts.view(key_pts.size(0), -1)
key_pts = key_pts.type(torch.FloatTensor).to(device)
images = images.type(torch.FloatTensor).to(device)
output_pts = net(images)
loss = criterion(output_pts, key_pts)
valid_loss += loss.item()*images.data.size(0)
train_loss = train_loss/len(train_loader.dataset)
valid_loss = valid_loss/len(test_loader.dataset)
print('Epoch: {} \tTraining Loss: {:.6f}'.format(epoch+1,train_loss,valid_loss))
if valid_loss <= valid_loss_min:
print('Validation loss decreased ({:.6f} --> {:.6f}). Saving model ...'.format(valid_loss_min,valid_loss))
torch.save(net,f'X:\\xxxx\\xxx\\xxx\\epoch{epoch + 1}_loss{valid_loss}.pth')
valid_loss_min = valid_loss
history['epoch'].append(epoch + 1)
history['train_loss'].append(train_loss)
history['valid_loss'].append(valid_loss)
print('Finished Training')
return history
'''
Above is the training code for reference!