# 获取测试集的加载器
test_loader = cfg.dataset_loader(root=cfg.mnist_dir,
train=False,
shuffle=False,
data_preprocess=test_data_preprocess)
# ---------------构建网络、定义损失函数、优化器--------------------------
# net = resnet18()
# net = resnet_v2.resnet18(num_classes=cfg.num_classes, type_dataset='cifar-10')
# net = vggnet.VGG(vgg_name='VGG11', num_classes=10, dataset='cifar-10')
net = LeNet(num_classes=10)
# 重写网络最后一层
# fc_in_features = net.fc.in_features # 网络最后一层的输入通道
# net.fc = nn.Linear(in_features=fc_in_features, out_features=cfg.num_classes)
# 加载模型权重、将网络放入GPU
if os.path.exists(cfg.checkpoints):
net.load_state_dict(torch.load(cfg.checkpoints))
print('load model argument...')
# -------------进行测试-----------------
print('进行测试.....')
# 测试函数
print('test stage...\n')
# 将网络结构调成验证模式、定义准确率、标签列表和预测列表
net.eval()
# cat_file = '/home/team/xiaonan/Dataset/cat_dog/test/cat/'
# dog_file = '/home/team/xiaonan/Dataset/cat_dog/test/'
#
# cat_list = os.listdir(cat_file)
# dog_list = os.listdir(dog_file)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
if args.resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.isdir(
'checkpoint'), 'Error: no checkpoint directory found!'
if os.path.isfile('./checkpoint/lifted_ckpt.t7'):
checkpoint = torch.load('./checkpoint/lifted_ckpt.t7')
best_acc = checkpoint['acc']
else:
checkpoint = torch.load('./checkpoint/classification_ckpt.t7')
best_acc = 0
net.load_state_dict(checkpoint['net'])
start_epoch = checkpoint['epoch']
current_best_epoch = start_epoch
print('the current best acc is %.3f on epoch %d' %
(best_acc, start_epoch))
criterion = LiftedStructureLoss()
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
for epoch in range(start_epoch, start_epoch + args.epochs):
train(epoch)
test(epoch)
def main(args):
check_path(args)
# CIFAR-10的全部类别，一共10类
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
# 数据集
data_builder = DataBuilder(args)
dataSet = DataSet(data_builder.train_builder(),
data_builder.test_builder(), classes)
# 选择模型
if args.lenet:
net = LeNet()
model_name = args.name_le
elif args.vgg:
net = Vgg16_Net()
model_name = args.name_vgg
elif args.resnet18:
net = ResNet18()
model_name = args.name_res18
elif args.resnet34:
net = ResNet34()
model_name = args.name_res34
elif args.resnet50:
net = ResNet50()
model_name = args.name_res50
elif args.resnet101:
net = ResNet101()
model_name = args.name_res101
elif args.resnet152:
net = ResNet152()
model_name = args.name_res152
# 交叉熵损失函数
criterion = nn.CrossEntropyLoss()
# SGD优化器
optimizer = optim.SGD(net.parameters(),
lr=args.learning_rate,
momentum=args.sgd_momentum,
weight_decay=args.weight_decay)
# 余弦退火调整学习率
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=150)
# 模型的参数保存路径
model_path = os.path.join(args.model_path, model_name)
# 启动训练
if args.do_train:
print("Training...")
trainer = Trainer(net, criterion, optimizer, scheduler,
dataSet.train_loader, dataSet.test_loader,
model_path, args)
trainer.train(epochs=args.epoch)
# t.save(net.state_dict(), model_path)
# 启动测试，如果--do_train也出现，则用刚刚训练的模型进行测试
# 否则就使用已保存的模型进行测试
if args.do_eval:
if not args.do_train and not os.path.exists(model_path):
print(
"Sorry, there's no saved model yet, you need to train first.")
return
# --do_eval
if not args.do_train:
checkpoint = t.load(model_path)
net.load_state_dict(checkpoint['net'])
accuracy = checkpoint['acc']
epoch = checkpoint['epoch']
print("Using saved model, accuracy : %f epoch: %d" %
(accuracy, epoch))
tester = Tester(dataSet.test_loader, net, args)
tester.test()
if args.show_model:
if not os.path.exists(model_path):
print(
"Sorry, there's no saved model yet, you need to train first.")
return
show_model(args)
if args.do_predict:
device = t.device("cuda" if t.cuda.is_available() else "cpu")
checkpoint = t.load(model_path, map_location=device)
net.load_state_dict(checkpoint['net'])
predictor = Predictor(net, classes)
img_path = 'test'
img_name = [os.path.join(img_path, x) for x in os.listdir(img_path)]
for img in img_name:
predictor.predict(img)
