def test_main():
print('-------test-----------')
gallery_path = images['gallery'].imgs
query_path = images['query'].imgs
gallery_cam, gallery_label = get_id(gallery_path)
query_cam, query_label = get_id(query_path)
model_structure = PCB(751)
if args.RPP:
model_structure = model_structure.convert_to_rpp()
model = load_network(model_structure)
model = PCB_test(model, feature_H)
model = model.eval()
if use_gpu:
model = model.cuda()
gallery_feature = extract_feature(model, datas['gallery'])
query_feature = extract_feature(model, datas['query'])
result = {
'gallery_feature': gallery_feature.numpy(),
'gallery_label': gallery_label,
'gallery_cam': gallery_cam,
'query_feature': query_feature.numpy(),
'query_label': query_label,
'query_cam': query_cam
}
if not os.path.isdir(args.result_dir):
os.mkdir(args.result_dir)
if args.RPP:
save_pre = '/RPP_'
else:
save_pre = '/PCB_'
if args.feature_H:
save_pre += 'H_'
else:
save_pre += 'G_'
scipy.io.savemat(args.result_dir + save_pre + 'result.mat', result)
return camera_id, labels
gallery_path = image_datasets['gallery'].imgs
query_path = image_datasets['query'].imgs
gallery_cam, gallery_label = get_id(gallery_path)
query_cam, query_label = get_id(query_path)
######################################################################
# Load Collected data Trained model
print('-------test-----------')
model_structure = PCB(751)
if args.RPP:
model_structure = model_structure.convert_to_rpp()
model = load_network(model_structure)
model = PCB_test(model, feature_H)
# Change to test mode
model = model.eval()
if use_gpu:
model = model.cuda()
# Extract feature
gallery_feature = extract_feature(model, dataloaders['gallery'])
query_feature = extract_feature(model, dataloaders['query'])
# Save to Matlab for check
#---------------------------------5、保存训练的模型------------------------------------------
def save_network(network, epoch_label):
save_filename = 'net_%s.pth' % epoch_label
save_path = os.path.join('/content/GPUID/model', name, save_filename)
torch.save(network.cpu().state_dict(), save_path) #保存为CPU版本的模型参数
if torch.cuda.is_available():
network.cuda(gpu_ids[0]) #模型回到gpu
#---------------------------------6、设置网络、加载模型--------------------------------------
#这里加入了len(classname)，可以根据数据集判断多少类，从而适用不同的数据集
if opt.RPP:
model = PCB(len(class_names))
model = model.convert_to_rpp()
opt.nclasses = len(class_names) #记录分类数
print(model) #输出模型结构
#ignored_params = list(map(id, model.avgpool.parameters()))
#base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
#optimizer_ft = optim.SGD([
#{'params': base_params, 'lr': 0.00},
#{'params': model.avgpool.parameters(), 'lr': 0.01},
#], weight_decay=5e-4, momentum=0.9, nesterov=True)
ignored_params = list(map(id, model.model.fc.parameters()))
ignored_params += (list(map(id, model.avgpool.parameters())) +
list(map(id, model.classifier0.parameters())) +
list(map(id, model.classifier1.parameters())) +
list(map(id, model.classifier2.parameters())) +
list(map(id, model.classifier3.parameters())) +
