Pytorch无法保存Torchvision模型的预训练重量

时间:2019-11-22 02:36:48

标签: save pytorch

对于某些分类任务,我尝试微调Torchvision模型。当我训练初始代码并保存最佳精度模型时,发生了一些奇怪的事情。这是简化的训练和测试代码:

net = torchvision.models.inception_v3(pretrained=True)
for epoch in range(60):
    for data in trainLoader:
         ....
    for data in testLoader:
         ....
    if test_acc > best_acc:
       best_acc = test_acc
       best_model_state_dict = copy.deepcopy(net.state_dict())
       best_epoch = epoch
       state = {
            "epoch": best_epoch,
             "model_state": best_model_state_dict,
             "best_acc": best_acc
            }

       save_path = "inception_model/{}_inceptionv3_{}_classroomClassification.pkl".format(str(best_acc), exp_i)
       torch.save(state, save_path)

当我使用下面的代码分别测试相同的数据时,准确性低于记录的值。

net = torchvision.models.inception_v3()
net.aux_logits = False
net.fc = nn.Linear(2048, 6)
model_path = r"D:\aha\0.8785714285714286_inceptionv3_2_classroomClassification.pkl"

state = torch.load(model_path)
net.load_state_dict(state["model_state"])
net.eval().cpu()
conf_mat = np.zeros([6, 6])
test_running_loss = 0.0
test_acc = 0.0
for data in tqdm(testLoader):
    img_feature, label = data
    img_feature = Variable(img_feature).cpu()
    net.eval()
    img_feature = Variable(img_feature).cpu()
    label = Variable(label).long().cpu()
    out = net(img_feature)
    pred = torch.max(out.data, 1)[1]
    test_acc += torch.sum(pred == label).item()
test_acc = test_acc / len(testData)
print(test_acc)

但是,当我将net = torchvision.models.inception_v3()更改为net = torchvision.models.inception_v3(pretrained=True)时,精度又变回来了。似乎pytorch无法保存预训练的权重。

0 个答案:

没有答案