Pytorch loading saved weights but not working - computer-vision

Question

I have this code. I save weights after every epoch and the code saves it. But when I load the weights the loss value starts from the initial loss value which means the loading is failing somehow.
net = torchvision.models.resnet18(pretrained=True)
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs, 136)
def train():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
criterion = L1Loss(reduction='sum')
lr = 0.0000001
optimizer = optim.Adam(net.parameters(), lr=lr, weight_decay=0.0005)
net.to(device)
state = torch.load('face2.txt')
net.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])
for epoch in range(int(0), 200000):
for batch, data in enumerate(trainloader, 0):
torch.cuda.empty_cache()
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = net(inputs).reshape(-1, 68, 2)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
state = {
'epoch': epoch,
'state_dict': net.state_dict(),
'optimizer': optimizer.state_dict(),
}
torch.save(state, 'face2.txt')
if __name__ == '__main__':
train()
The initial loss is 50k plus and after some thousand epochs the loss becomes 50-60. Now when I re-run the code I expect it to start from a near loss value but it starts from around 50k again.

Answer 1

The code you wrote:
net = torchvision.models.resnet18(pretrained=True)
means you start over again with the same net - pretrained resnet18. Instead, you should load the last state if present, this would fix your problem.
I will slightly update your notation:
state = {
'epoch': epoch,
'model_state_dict': net.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
...
}
torch.save(state, filepath)
Learnable parameters are the first state_dict (model state dict).
The second state_dict is the optimizer state dict. You recall that the optimizer is used to improve our learnable parameters. But the optimizer state_dict is fixed. Nothing to learn in there.
Your code at some point should look like:
model.load_state_dict(state['model_state_dict'])
optimizer.load_state_dict(state['optimizer_state_dict'])