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())) +