def main():
# you should set data path on the top
global train_path, val_path, test_path
global train_batch_size
if 'Flower102' in args.data_name:
train_path = "./Flower102/train"
test_path = "./Flower102/test"
val_path = "./Flower102/val"
elif 'Birds200' in args.data_name:
train_path = "./Birds200/train"
test_path = "./Birds200/test"
# global train_path, test_path
if 'Flower102' in train_path:
model = ModifiedAlexNet(102)
train_loader = dataset.train_loader(train_path,
batch_size=train_batch_size,
num_workers=4,
pin_memory=True)
val_loader = dataset.test_loader(val_path,
batch_size=1,
num_workers=4,
pin_memory=True)
test_loader = dataset.test_loader(test_path,
batch_size=1,
num_workers=4,
pin_memory=True)
elif 'Birds200' in train_path:
model = ModifiedAlexNet(200)
train_loader = dataset.train_loader(train_path,
batch_size=train_batch_size,
num_workers=4,
pin_memory=True)
test_loader = dataset.test_loader(test_path,
batch_size=1,
num_workers=4,
pin_memory=True)
if use_gpu:
model = model.cuda()
print("Use GPU!")
else:
print("Use CPU!")
checkpoint = torch.load('./models/alexnet_Flower102_0.787.pth')
model.load_state_dict(checkpoint, strict=True)
if 'Flower102' in train_path:
train_val_test(model,
train_loader,
val_loader,
test_loader,
optimizer=None,
epoches=50)
elif 'Birds200' in train_path:
train_test(model,
train_loader,
test_loader,
optimizer=None,
epoches=10)
def main():
global args, num_batches, use_gpu
if not args.zeroTrain:
if args.arch == "densenet201":
model = models.Modified_densenet201()
elif args.arch == "Resnet50":
model = models.Modified_Resnet50()
elif args.arch == "Resnet101":
model = models.Modified_Resnet101()
elif args.arch == "Resnet152":
model = models.Modified_Resnet152()
elif args.arch == "nasnetalarge":
model = models.Modified_nasnetalarge()
if args.resume:
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint, strict=True)
if use_gpu:
model = model.cuda()
print("Use GPU!")
else:
print("Use CPU!")
train_path = args.data_path[0]
val_path = args.data_path[1]
test_path = args.data_path[2]
train_loader = dataset.train_loader(train_path,
batch_size=args.batch_size,
num_workers=10,
pin_memory=True)
val_loader = dataset.test_loader(val_path,
batch_size=1,
num_workers=10,
pin_memory=True)
test_loader = dataset.test_loader(test_path,
batch_size=1,
num_workers=10,
pin_memory=True)
if args.TenCrop:
TenCroptest_loader = dataset.testTenCrop_loader(test_path,
batch_size=1,
num_workers=10,
pin_memory=True)
train_val_test(model,
train_loader,
val_loader,
test_loader,
print_freq=50,
TenCroptest_loader=TenCroptest_loader,
optimizer=None,
epoches=args.epochs)
def finetune():
global infoLogger
global use_gpu, num_batches, train_batch_size
if not zero_train:
global train_path, test_path
global train_batch_size
if 'Flower102' in data_name:
train_path = "./Flower102/train"
test_path = "./Flower102/test"
num_classes = 102
elif 'Birds200' in data_name:
train_path = "./Birds200/train"
test_path = "./Birds200/test"
num_classes = 200
elif 'catdog' in data_name:
train_path = "./CatDog/train"
test_path = "./CatDog/test"
num_classes = 2
if not checkpoint:
if use_gpu:
print("Use GPU!")
model = firstreloadParam(num_classes)
else:
checkpointPath = './1x1models/finetune/alexnet1x1_Flower102_0.821.pth'
state = torch.load(checkpointPath)
if use_gpu:
print("Use GPU!")
model = add1x1layers(num_classes).cuda()
else:
print("Use CPU!")
model = add1x1layers(num_classes)
model.load_state_dict(state)
train_loader = dataset.train_loader(train_path,
batch_size=train_batch_size,
num_workers=4,
pin_memory=True)
test_loader = dataset.test_loader(test_path,
batch_size=1,
num_workers=4,
pin_memory=True)
infoLogger.info("dataset is: " + args.data_name)
optimizer = SGD(
[{
'params': model.classifier.parameters()
} # 浅层1x1不好剪枝, 先冻住
],
weight_decay2=5e-4,
lr=1e-3,
momentum=0.9)
infoLogger.info("weight_decay2 is: " + str(5e-4))
infoLogger.info("learning rate is: " + str(1e-3))
train_test(model,
train_loader,
test_loader,
optimizer=optimizer,
epoches=20)
def main():
time1 = time()
global args, use_gpu
input_size = args.input_size
resume1 = './checkpoint/BBX/NewAnnoCleaned/Resnet152_64_{}_Epoch.pth'
# resume2 = './checkpoint/BBX/sigmoid_model/Lr_0_001/Five_Branch/Bottleneck/Five_Depth/Resnet152_64_{}_Epoch.pth'
test_loader = dataset.test_loader(args.test_path,
args.relative_path,
batch_size=args.batch_size,
num_workers=10,
pin_memory=True)
model = models.sigmoid_ResnetModel(args.arch, 224, args.num_classs, 2, 17,
5)
for i in range(49, 48, -1):
f = open(
'./result/NewAnnoCleaned/LR_0_001_FiveDepth_epoch_{}_top3_top5.txt'
.format(str(i)), 'w')
print("{} epoch".format(str(i)))
checkpoint = torch.load(resume1.format(str(i)))
model.load_state_dict(checkpoint, strict=True)
f.write("{} epoch\n".format(str(i)))
if use_gpu:
model = model.cuda()
else:
pass
if args.Branch == 1:
# _, _, overall_top3_recall, overall_top5_recall = top3_5_recall(model, test_loader, topk=(3,5), epoch=i)
pass
elif args.Branch == 17:
_, _, overall_top3_recall, overall_top5_recall = recall_precision(
model, test_loader, topk=(3, 5), epoch=i)
f.write("top-3 ")
for key in overall_top3_recall.keys():
f.write("{} ".format(key))
f.write("\n")
for _, value in overall_top3_recall.items():
f.write("{} ".format(str(round(value * 100, 3))))
f.write("\n")
f.write("top-5 ")
for key in overall_top5_recall.keys():
f.write("{} ".format(key))
f.write("\n")
for _, value in overall_top5_recall.items():
f.write("{} ".format(str(round(value * 100, 3))))
f.write("\n")
f.close()
print("total time is: {} seconds".format(round(time() - time1, 2)))
def main():
global args, use_gpu
input_size = args.input_size
model = models.Modified_Resnet(args.arch, 46, input_size)
if use_gpu:
model = model.cuda()
print("Use GPU!")
else:
print("Use CPU!")
test_root = args.test_path
test_loader = dataset.test_loader(test_root,
args.input_size,
batch_size=args.batch_size,
num_workers=10,
pin_memory=True)
f = open('./top1_3_5.txt', 'w')
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint, strict=True)
top1_3_5(f, 49, model, test_loader)
f.close()
num_classes = 102
elif 'Birds200' in data_name:
train_path = "./Birds200/train"
test_path = "./Birds200/test"
num_classes = 200
elif 'catdog' in data_name:
train_path = "./CatDog/train"
test_path = "./CatDog/test"
num_classes = 2
train_loader = dataset.train_loader(train_path,
batch_size=train_batch_size,
num_workers=4,
pin_memory=True)
test_loader = dataset.test_loader(test_path,
batch_size=1,
num_workers=4,
pin_memory=True)
infoLogger.info("dataset is: " + args.data_name)
def test(model, test_loader):
model.eval()
test_correct = 0
test_total = 0
for i, (batch, label) in enumerate(test_loader):
batch = batch.cuda()
output = model(Variable(batch))
pred_label = output.data.max(1)[1] # 返回模型预测概率最大的标签
test_correct += pred_label.cpu().eq(
label).sum() # label为torch.LongTensor类型
def finetune():
global train_batch_size
if arch == "mobilenetv1":
model = models.MobileNet(20)
if use_gpu:
model = model.cuda()
elif arch == "mobilenetv2":
model = models.MobileNetV2(20)
if use_gpu:
model = model.cuda()
elif arch == 'squeezenet1_1':
model = models.squeezenet1_1(20, pretrained=False)
if use_gpu:
model = model.cuda()
if checkpoint:
if arch == "mobilenetv2":
checkpointPath = './checkpoint/mobilenetv20.865.pth'
state = torch.load(checkpointPath)
model.load_state_dict(state)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
if arch == 'squeezenet1_1':
checkpointPath = './checkpoint/squeezenet1_10.762.pth'
state = torch.load(checkpointPath)
model.load_state_dict(state)
optimizer = optim.SGD(model.newclassifier.parameters(),
lr=0.01,
momentum=0.9)
else:
if arch == "mobilenetv1":
state = torch.load(
'./checkpoint/mobilenet_sgd_rmsprop_69.526.tar')['state_dict']
model.load_state_dict(state, strict=False)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
elif arch == "mobilenetv2":
state = torch.load(
'./checkpoin t/mobilenetv2_Top1_71.806_Top2_90.410.pth.tar')
model.load_state_dict(state, strict=False)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
elif arch == 'squeezenet1_1':
model = models.squeezenet1_1(20, pretrained=True)
if use_gpu:
model = model.cuda()
optimizer = optim.SGD(model.newclassifier.parameters(),
lr=0.01,
momentum=0.9)
# train_path = '/home/smiles/ModelPruning/CatDog/train'
# test_path = '/home/smiles/ModelPruning/CatDog/test'
train_path = './train'
# test_path = train_path
test_path = './test'
train_loader = dataset.train_loader(train_path,
batch_size=train_batch_size,
num_workers=4,
pin_memory=True)
test_loader = dataset.test_loader(test_path,
batch_size=1,
num_workers=4,
pin_memory=True)
train_test('./checkpoint/',
arch,
model,
train_loader,
test_loader,
optimizer=optimizer,
epoches=40)