3.2 model saving and loading

Table of Contents

Model preservation

General model loading 　

Save model details

Model parameter freezing

---

Model preservation

#  Save the entire network 

torch.save(net, PATH)

#  Save parameters in the network ,  Fast , Less space 

torch.save(net.state_dict(),PATH)

General model loading 　

model_dict=torch.load(PATH)

model_dict=model.load_state_dict(torch.load(PATH))

Save model details

torch.save({'epoch': epochID + 1, 'state_dict': model.state_dict(), 'best_loss': lossMIN,

'optimizer': optimizer.state_dict(),'alpha': loss.alpha, 'gamma': loss.gamma},

checkpoint_path + '/m-' + launchTimestamp + '-' + str("%.4f" % lossMIN) + '.pth.tar')

ｅｇ：    
for epoch in range(init_epoch, basic_configs['num_epochs']):
        logger.info("Begin training epoch {}".format(epoch + 1))
        train_function(epoch)

        net_checkpoint_name = args.exp + "_net_epoch" + str(epoch + 1)
        net_checkpoint_path = os.path.join(exp_ckpt_dir, net_checkpoint_name)
        net_state = {"epoch": epoch + 1,
                     "network": net.module.state_dict()}
        torch.save(net_state, net_checkpoint_path)

Model parameter freezing  

for p in self.parameters():
    p.requires_grad = False


class RESNET_YYT(nn.Module):

    def init(self, model, pretrained):

        super(RESNET_YYT, self).__init__()

        self.resnet = model(pretrained)

        for p in self.parameters():

            p.requires_grad = False

        self.f = SpectralNorm(nn.Conv2d(2048, 512, 1))

        self.g = SpectralNorm(nn.Conv2d(2048, 512, 1))

        self.h = SpectralNorm(nn.Conv2d(2048, 2048, 1))

Print network layers

model_dict = torch.load('/001_net_epoch5loss0.08947')['network'] 
dict_name = list(model_dict)
for i, p in enumerate(dict_name):
    print(i, p)