pytorch如何保存训练模型参数并实现继续训练

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(64, 32)
        self.linear1 = nn.Linear(32, 10)
        self.relu = nn.ReLU()
    def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)
        x = self.linear1(x)
        return x

    rand1 = torch.rand((100, 64)).to(torch.float)
    label1 = np.random.randint(0, 10, size=100)
    label1 = torch.from_numpy(label1).to(torch.long)
    rand2 = torch.rand((100, 64)).to(torch.float)
    label2 = np.random.randint(0, 10, size=100)
    label2 = torch.from_numpy(label2).to(torch.long)

    model = MyModel()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    loss = nn.CrossEntropyLoss()
    iteration = 100
    for i in range(iteration):
        output = model(rand1)
        my_loss = loss(output, label1)
        optimizer.zero_grad()
        my_loss.backward()
        optimizer.step()
        print("iteration:{} loss:{}".format(i, my_loss))

torch.save(model, save_path)

def save_model(save_path, iteration, optimizer, model):
    torch.save({'iteration': iteration,
                'optimizer_dict': optimizer.state_dict(),
                'model_dict': model.state_dict()},
                save_path)
    print("model save success")

def load_model(save_name, optimizer, model):
    model_data = torch.load(save_name)
    model.load_state_dict(model_data['model_dict'])
    optimizer.load_state_dict(model_data['optimizer_dict'])
    print("model load success")

    path = "net.dict"
    save_model(path, iteration, optimizer, model)
    print(model.state_dict()['linear.weight'])
    new_model = MyModel()
    new_optimizer = torch.optim.Adam(new_model.parameters(), lr=0.01)
    load_model(path, new_optimizer, new_model)
    print(new_model.state_dict()['linear.weight'])

    for i in range(iteration):
        output = new_model(rand2)
        my_loss = loss(output, label2)
        new_optimizer.zero_grad()
        my_loss.backward()
        new_optimizer.step()
        print("iteration:{} loss:{}".format(i, my_loss))

# 存储时
'scheduler': scheduler.state_dict()
# 读取时
scheduler.load_state_dict(checkpoint['lr_schedule'])