Hello,
for the last 2 days I am trying to solve issue when resuming training from model checkpoint. Problem is that the training loss after resuming is a LOT different than before saving model (the difference is huge, almost as if the model was right after initialization process). I can see, that after few iterations it increases accuracy (decreases loss) much faster than if it was learned from scratch. But I don’t know what cuases the loss to be so high.
But weird thing is that I checked output of the same input (before I turned models in the eval mode - model.eval()) right before saving model and right after loading and the output was the same -> so that means it has correct weights everywhere in the model. I assume that the model has to be loaded correctly, because otherwise it would produce different output for the same inputs.
So I was wondering if the optimizer is saved and loaded correctly, so as the lr_scheduler. I checked and optimizer is loaded correctly, the only difference was in the ids for each layer, but I guess these ids are picked randomly at time the model is created. But the order was correct and the values in the state dict were also correct. The lr_scheduler was having correct values. Also I tried the model without lr_scheduler and removed all the “params” from optimizer (I’ve set momentum = 0, weight_decay = 0, so there’s actually no params in the optimizer). But still without success.
Maybe there can be some problem with using optimizer with 2 models like this optimizer = torch.optim.SGD([{'params': model.parameters()}, {'params': fc_layer.parameters()}], lr=config.lr, weight_decay=config.weight_decay, momentum=config.momentum) ? I am training face recognition model with AM-Softmax (see the code below), so that’s why I have 2 models (one model for feature-vector extraction, second for learning using AM-Softmax). Can be the problem here? I checked and the weights are correctly saving and loading from and to the fc_layer model.
class AddMarginProduct(nn.Module):
def __init__(self, in_features: int, out_features: int, s: float = 30.0, m: float = 0.35) -> None:
'''
Implementation of large margin cosine distance
Arguments:
nn {[type]} -- [description]
in_features {int} -- [size of each input sample]
out_features {int} -- [size of each output sample]
Keyword Arguments:
s {float} -- [norm of input feature] (default: {30.0})
m {float} -- [margin] (default: {0.35})
'''
super(AddMarginProduct, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.s = s
self.m = m
self.weight = Parameter(torch.FloatTensor(out_features, in_features))
nn.init.xavier_uniform_(self.weight)
def forward(self, input, label):
'''
Forward pass of layer
Arguments:
input {[type]} -- [torch layer inputs]
label {[type]} -- [torch model labels]
Returns:
[type] -- [output]
'''
# --------------------------- cos(theta) & phi(theta) ---------------------------
cosine = F.linear(F.normalize(input), F.normalize(self.weight))
phi = cosine - self.m
# --------------------------- convert label to one-hot ---------------------------
one_hot = torch.zeros(cosine.size()).scatter_(1, label.unsqueeze(1), 1).byte()
output = torch.where(one_hot, phi, cosine)
output *= self.s
return output
Now I am out of ideas why there’s the big difference after resuming from the model checkpoint.
There’s just main parts of my code, if you can spot some mistake I did? (I removed unnecessary part as config parsing, my custom classes for progress managing and so on, but the complete code doesn’t yield any error).
def main():
# Set cuda if it's available, set correct default tensor and set random seed
use_cuda = not config.no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
# torch.manual_seed(config.seed) # Tried with or without manual seed
if device.type == 'cuda':
torch.set_default_tensor_type('torch.cuda.FloatTensor')
cudnn.benchmark = True
# Create train dataset factory with transforms and train_loader
train_dataset = dataset_factory.get_train_dataset(transforms=T.Compose([
T.RandomAffine(10, translate=(0.05, 0.05), scale=(0.9, 1.1), shear=5),
T.RandomHorizontalFlip(p=0.5),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]))
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.num_of_workers,
pin_memory=config.pin_memory,
)
# Create test dataset with transforms and test_loader
test_dataset = dataset_factory.get_test_dataset(transforms=T.Compose([
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]))
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.test_batch_size,
shuffle=False,
num_workers=config.num_of_workers,
pin_memory=config.pin_memory,
)
# Create models
model = resnet_18(pretrained=config.pretrained, num_classes=config.feature_size)
fc_layer = AddMarginProduct(config.feature_size, train_dataset.one_hot_vector_classes)
# Cast to device
model.to(device)
fc_layer.to(device)
# Create loss func and optimizer
criterion = F.cross_entropy
optimizer = torch.optim.SGD([{'params': model.parameters()}, {'params': fc_layer.parameters()}], lr=config.lr, weight_decay=config.weight_decay, momentum=config.momentum)
# Set scheduler and set initial epoch
scheduler = StepLR(optimizer, step_size=config.lr_step_size, gamma=config.lr_gammma)
start_epoch = 1
# Model manager is just wrapper for torch.save() and torch.load(), because it automatically adds best accuracies to the dictionary
model_manager = ModelManager()
# If we should resume
if args.resume and os.path.isfile(args.resume):
fc_load_failed = False
# Load the checkpoint
print('Loading checkpoint {}'.format(args.resume))
checkpoint_state = model_manager.load_model_checkpoint(args.resume)
# Models
model.load_state_dict(checkpoint_state["models"]["model"])
try:
fc_layer.load_state_dict(checkpoint_state["models"]["fc_layer"])
# fc_layer.to(device)
except Exception as e:
fc_load_failed = True
print("Cannot load fc_layer (probably wrong dimensions from loaded model) - exception: ", e)
# Load model only (not the optimizer and other options)
if not args.model_only:
scheduler.load_state_dict(checkpoint_state['scheduler'])
if not fc_load_failed:
optimizer.load_state_dict(checkpoint_state['optimizer'])
start_epoch = checkpoint_state['epoch'] + 1
config.epochs += start_epoch
print("Loaded checkpoint '{}' (epoch {})".format(args.resume, start_epoch))
# Train
for epoch in range(start_epoch, config.epochs):
scheduler.step()
# In this function is called model.train() and fc_layer.train()
train(config, model, fc_layer, device, train_loader, optimizer, criterion, epoch, update_logger, epoch_logger)
# There is called model.eval() and fc_layer.eval()
lfw_acc = test(config, model, fc_layer, device, test_loader, criterion, epoch_logger, epoch, evaluator, plots_dir)
# Save model after each epoch
model_manager.save_model_checkpoint(
os.path.join(models_dir, "{}_e_{}_acc_{}.pth".format(config.model, str(epoch).zfill(4), lfw_acc)),
{
"model_arch": config.model,
"fc_layer_arch": config.fc_layer,
"models": {
"model": model.state_dict(),
"fc_layer": fc_layer.state_dict()
},
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
"epoch": epoch,
},
lfw_acc,
os.path.join(models_dir, "{}_best.pth".format(config.model))
)
Have I done anything wrong?
Also I tried pytorch version 0.4.1 and 1.0.0, and nothing changed.
Thank you very much for your help, I am a little desperate, because I am out of ideas what could be wrong.