# 记录下验证集精度
if val_loss < best_loss:
best_epoch, best_loss = epoch, val_loss
# 保存model可学习参数
torch.save(model.state_dict(), 'Model/model.pt')
# 打印相关信息
Tools.printInfo(epoch, train_loss, val_loss, best_epoch, best_loss)
else:
# 预测过程
test_path, test_label = Tools.dataFromPath(test_img_path)
test_dataset = SVHNDataset(
test_path, test_label,
transforms.Compose([
transforms.Resize((64, 128)),
transforms.ColorJitter(0.3, 0.3, 0.2),
transforms.RandomRotation(5),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
]))
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=5,
)
model = SVHN_Model1()
model.load_state_dict(torch.load("Model/model.pt", map_location='cpu'))
pred_labels = model.myPredicting(test_loader)
Tools.submit(demo_submit_path, pred_labels)
