def main_train(): f = open(args.save_dir + 'train_log.txt', 'w') copyfile('./train.py', args.save_dir + '/train.py') copyfile('./model.py', args.save_dir + '/model.py') model = PCB(len(datas["class"])) if gpu: model = model.cuda() if is_parallel_train: model = nn.DataParallel(model, device_ids=gpu_ids) criterion = nn.CrossEntropyLoss() model = pcb_train(model, criterion, f, "PCB", 60) if args.RPP: model = get_net(is_parallel_train, model).convert_to_rpp() if use_gpu: model = model.cuda() if is_parallel_train: model = nn.DataParallel(model, device_ids=gpu_ids) model = rpp_train(model, criterion, f, "RPP", 5) model = full_train(model, criterion, f, "full", 10) f.close()
def get_reid(model_path): checkpoint = torch.load(model_path) model_structure = PCB(1453) model_structure.load_state_dict(checkpoint) model = PCB_test(model_structure) model = model.eval() return model
def main(): transform_test = T.Compose([ T.Resize((288,144),interpolation=3), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) print("I am in the main") use_gpu = torch.cuda.is_available() #model = torch.load(os.path.join('./model','best_model.pth.tar')) #print(type(model)) #model.eval() #model = models.init_model(name=args.arch, num_classes=751, loss={'xent'}) #if use_gpu: # model = nn.DataParallel(model).cuda() #checkpoint = torch.load(os.path.join('./model','net_last.pth')) #save_path = os.path.join('./model','net_last.pth') #model.load_state_dict(checkpoint['state_dict']) #model.load_state_dict(torch.load(save_path)) #model = torch.load(os.path.join('./model','net_last.pth')) #if use_gpu: # model = nn.DataParallel(model).cuda() #model.to(device) #model_structure = models.init_model(name=args.arch, num_classes=751, loss={'xent'}) model_structure = PCB(751) model=load_network(model_structure) model = PCB_test(model) model = model.eval() if use_gpu: model = model.cuda() for subset in ['query', 'gallery']: print(subset) test_names, test_features = extractor(model, DataLoader(Dataset(subset,transform=transform_test))) results = {'names': test_names, 'features': test_features.numpy()} scipy.io.savemat(os.path.join('example', 'feature_val_%s.mat' % (subset)), results)
def load_model(self): if self.use_dense: model_structure = ft_net_dense(self.nclasses) elif self.use_NAS: model_structure = ft_net_NAS(self.nclasses) else: model_structure = ft_net(self.nclasses, stride=self.stride) if self.PCB: model_structure = PCB(self.nclasses) #if self.fp16: # model_structure = network_to_half(model_structure) self.model = self.load_network(model_structure) # Remove the final fc layer and classifier layer if self.PCB: #if self.fp16: # model = PCB_test(model[1]) #else: self.model = PCB_test(self.model) else: #if self.fp16: #model[1].model.fc = nn.Sequential() #model[1].classifier = nn.Sequential() #else: self.model.classifier.classifier = nn.Sequential() # Change to test mode self.model = self.model.eval() if self.use_gpu: self.model = self.model.cuda()
def load_network(name, finetune=False): # Load config if finetune: dirname = os.path.join('./model', 'ft_ResNet') else: dirname = os.path.join('./model', name) last_model_name = os.path.basename(get_model_list(dirname, 'net')) print(last_model_name) epoch = last_model_name.split('_')[1] epoch = epoch.split('.')[0] if not epoch == 'last': epoch = int(epoch) if opt.PCB: model = PCB(len(class_names), return_f=True) else: model = ft_net(len(class_names), return_f=True) # load model if isinstance(epoch, int): save_filename = 'net_%03d.pth' % epoch else: save_filename = 'net_%s.pth' % epoch if finetune: save_path = os.path.join('./model', 'ft_ResNet', save_filename) else: save_path = os.path.join('./model', name, save_filename) print('Load the model from %s' % save_path) network = model network.load_state_dict(torch.load(save_path)) return network, epoch
def __init__(self): #base_model = VGG16(weights='imagenet') #self.model = Model(inputs=base_model.input, outputs=base_model.get_layer('fc1').output) save_path = os.path.join('./model', 'net.pth') model = PCB(575) self.model = PCB_test(model, num_parts=6, cluster_plots=False) self.model.load_state_dict(torch.load(save_path), strict=False) self.avgpool = nn.AdaptiveAvgPool2d((3, 2)) self.model = self.model.eval()
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)
def load_network(name, opt): # Load config dirname = os.path.join('./model', name) last_model_name = os.path.basename(get_model_list(dirname, 'net')) epoch = last_model_name.split('_')[1] epoch = epoch.split('.')[0] if not epoch == 'last': epoch = int(epoch) config_path = os.path.join(dirname, 'opts.yaml') with open(config_path, 'r') as stream: config = yaml.load(stream) opt.name = config['name'] opt.data_dir = config['data_dir'] opt.train_all = config['train_all'] opt.droprate = config['droprate'] opt.color_jitter = config['color_jitter'] opt.batchsize = config['batchsize'] opt.h = config['h'] opt.w = config['w'] opt.stride = config['stride'] if 'pool' in config: opt.pool = config['pool'] if 'h' in config: opt.h = config['h'] opt.w = config['w'] if 'gpu_ids' in config: opt.gpu_ids = config['gpu_ids'] opt.erasing_p = config['erasing_p'] opt.lr = config['lr'] opt.nclasses = config['nclasses'] opt.erasing_p = config['erasing_p'] opt.use_dense = config['use_dense'] opt.PCB = config['PCB'] opt.fp16 = config['fp16'] if opt.use_dense: model = ft_net_dense(opt.nclasses, opt.droprate, opt.stride, None, opt.pool) else: model = ft_net(opt.nclasses, opt.droprate, opt.stride, None, opt.pool) if opt.PCB: model = PCB(opt.nclasses) # load model if isinstance(epoch, int): save_filename = 'net_%03d.pth' % epoch else: save_filename = 'net_%s.pth' % epoch save_path = os.path.join('./model', name, save_filename) print('Load the model from %s' % save_path) network = model network.load_state_dict(torch.load(save_path)) return network, opt, epoch
def get_model(class_names_size): if opt.use_dense: model = ft_net_dense(class_names_size, opt.droprate) elif opt.use_NAS: model = ft_net_NAS(class_names_size, opt.droprate) else: model = ft_net(class_names_size, opt.droprate, opt.stride) if opt.PCB: model = PCB(class_names_size) return model
def loadModel(modelName, config, stateDictPath): nclasses = config["nclasses"] if modelName == "densenet121": model_structure = ft_net_dense(nclasses) elif modelName == "NAS": model_structure = ft_net_NAS(nclasses) elif modelName == "ResNet50": model_structure = ft_net(nclasses, stride=config["stride"]) elif modelName == "PCB": model_structure = PCB(nclasses) else: raise KeyError("Undefined modelName, %s" % modelName) return loadNetwork(model_structure, stateDictPath)
def main(): print("==========\nArgs:{}\n==========".format(opt)) train_data_dir = opt.train_data_dir test_data_dir = opt.test_data_dir name = opt.name os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_ids use_gpu = torch.cuda.is_available() if opt.use_cpu: use_gpu = False if not opt.evaluate: sys.stdout = Logger(osp.join(opt.save_dir, opt.train_log)) else: sys.stdout = Logger(osp.join(opt.save_dir, opt.test_log)) if use_gpu: print("Currently using GPU {}".format(opt.gpu_ids)) cudnn.benchmark = True else: print("Currently using CPU (GPU is highly recommended)") #str_ids = opt.gpu_ids.split(',') #gpu_ids = [] #for str_id in str_ids: # gid = int(str_id) # if gid >=0: # gpu_ids.append(gid) # ## set gpu ids #if len(gpu_ids)>0: # torch.cuda.set_device(gpu_ids[0]) #print(gpu_ids[0]) # Load train Data # --------- print("==========Preparing trian dataset========") transform_train_list = [ # transforms.RandomResizedCrop(size=128, scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC) transforms.Resize((288, 144), interpolation=3), transforms.RandomCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.PCB: transform_train_list = [ transforms.Resize((384, 192), interpolation=3), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.erasing_p > 0: transform_train_list = transform_train_list + [ RandomErasing(probability=opt.erasing_p, mean=[0.0, 0.0, 0.0]) ] if opt.color_jitter: transform_train_list = [ transforms.ColorJitter( brightness=0.1, contrast=0.1, saturation=0.1, hue=0) ] + transform_train_list train_data_transforms = transforms.Compose(transform_train_list) train_all = '' if opt.train_all: if opt.use_clean_imgs: train_all = '_all_clean' else: train_all = '_all' print("Using all the train images") train_image_datasets = {} train_image_datasets['train'] = datasets.ImageFolder( os.path.join(train_data_dir, 'train' + train_all), train_data_transforms) train_dataloaders = { x: torch.utils.data.DataLoader(train_image_datasets[x], batch_size=opt.train_batch, shuffle=True, num_workers=4) for x in ['train'] } dataset_sizes = {x: len(train_image_datasets[x]) for x in ['train']} class_names = train_image_datasets['train'].classes inputs, classes = next(iter(train_dataloaders['train'])) ###################################################################### # Prepare test data if not opt.train_only: print("========Preparing test dataset========") transform_test_list = transforms.Compose([ transforms.Resize((384, 192), interpolation=3), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) test_image_datasets = { x: datasets.ImageFolder(os.path.join(test_data_dir, x), transform_test_list) for x in ['gallery', 'query'] } test_dataloaders = { x: torch.utils.data.DataLoader(test_image_datasets[x], batch_size=opt.test_batch, shuffle=False, num_workers=4) for x in ['gallery', 'query'] } print("Initializing model...") if opt.use_dense: model = ft_net_dense(len(class_names)) else: model = ft_net(len(class_names)) if opt.PCB: model = PCB(len(class_names)) print("Model size: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) start_epoch = opt.start_epoch if opt.resume: print("Loading checkpoint from '{}'".format(opt.resume)) checkpoint = torch.load(opt.resume) model.load_state_dict(checkpoint['state_dict']) # model.load_state_dict(checkpoint) start_epoch = checkpoint['epoch'] if use_gpu: model = nn.DataParallel(model).cuda() if opt.evaluate: print("Evaluate only") test(model, test_image_datasets, test_dataloaders, use_gpu) return criterion = nn.CrossEntropyLoss().cuda() if opt.PCB: ignored_params = list(map(id, model.module.resnet50.fc.parameters())) ignored_params += ( list(map(id, model.module.classifier0.parameters())) + list(map(id, model.module.classifier1.parameters())) + list(map(id, model.module.classifier2.parameters())) + list(map(id, model.module.classifier3.parameters())) + list(map(id, model.module.classifier4.parameters())) + list(map(id, model.module.classifier5.parameters())) #+list(map(id, model.classifier7.parameters() )) ) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer = optim.SGD( [ { 'params': base_params, 'lr': 0.01 }, { 'params': model.module.resnet50.fc.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier0.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier1.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier2.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier3.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier4.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier5.parameters(), 'lr': 0.1 }, #{'params': model.classifier7.parameters(), 'lr': 0.01} ], weight_decay=5e-4, momentum=0.9, nesterov=True) else: ignored_params = list(map( id, model.module.model.fc.parameters())) + list( map(id, model.module.classifier.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer = optim.SGD([{ 'params': base_params, 'lr': 0.01 }, { 'params': model.module.model.fc.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier.parameters(), 'lr': 0.1 }], weight_decay=5e-4, momentum=0.9, nesterov=True) # Decay LR by a factor of 0.1 every 40 epochs if opt.stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.stepsize, gamma=0.1) start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") ###################################################################### # Training the model # ------------------ # # Now, let's write a general function to train a model. Here, we will # illustrate: # # - Scheduling the learning rate # - Saving the best model # # In the following, parameter ``scheduler`` is an LR scheduler object from # ``torch.optim.lr_scheduler``. for epoch in range(start_epoch, opt.max_epoch): start_train_time = time.time() if opt.train_only: print("==> Training only") train(epoch, model, criterion, optimizer, train_dataloaders, use_gpu) train_time += round(time.time() - start_train_time) if epoch % opt.eval_step == 0: if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'epoch': epoch, }, 0, osp.join(opt.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) if opt.stepsize > 0: scheduler.step() else: train(epoch, model, criterion, optimizer, train_dataloaders, use_gpu) train_time += round(time.time() - start_train_time) if opt.stepsize > 0: scheduler.step() if (epoch + 1) > opt.start_eval and opt.eval_step > 0 and ( epoch + 1) % opt.eval_step == 0 or (epoch + 1) == opt.max_epoch: print("==> Test") rank1 = test(model, test_image_datasets, test_dataloaders, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(opt.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) if not opt.train_only: print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
query_cam, query_label = get_id(query_path) if opt.multi: mquery_path = image_datasets['multi-query'].imgs mquery_cam, mquery_label = get_id(mquery_path) ###################################################################### # Load Collected data Trained model print('-------test-----------') if opt.use_dense: model_structure = ft_net_dense(751) else: model_structure = ft_net(751) if opt.PCB: model_structure = PCB(751) if opt.fp16: model_structure = network_to_half(model_structure) model = load_network(model_structure) # Remove the final fc layer and classifier layer if opt.PCB: model = PCB_test(model) model = network_to_half(model) elif opt.fp16: model[1].model.fc = nn.Sequential() model[1].classifier = nn.Sequential() else: model.model.fc = nn.Sequential()
gallery_path = image_datasets['gallery'].imgs query_path = image_datasets['query'].imgs g_path = [] for gp in gallery_path: (name, _) = gp g_path.append(name.replace(' ', '').split('/')[-1]) q_path = [] for qp in query_path: (name, _) = qp q_path.append(name.replace(' ', '').split('/')[-1]) ###################################################################### # Load Collected data Trained model print('-------test-----------') model_structure = PCB(class_num=4768) #for duke if torch.cuda.device_count() > 1: model_structure = nn.DataParallel(model_structure) model = load_network(model_structure) # 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
###################################################################### # Finetuning the convnet # ---------------------- # # Load a pretrainied model and reset final fully connected layer. # if opt.use_dense: model = ft_net_dense((id_class_number), opt.droprate) else: model = ft_net((id_class_number), attr_class_number, opt.droprate, opt.stride) if opt.PCB: model = PCB((id_class_number)) opt.nclasses = (id_class_number) print(model) if not opt.PCB: ignored_params = list(map(id, model.model.fc.parameters())) for i in range(attr_class_number + 1): list(map(id, model.__getattr__('class_' + str(i)).parameters())) # ignored_params = list(map(id, model.model.fc.parameters())) + \ # list(map(id, model.classifier.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer_ft = optim.SGD(
mquery_path = image_datasets['multi-query'].imgs mquery_cam, mquery_label = get_id(mquery_path) ###################################################################### # Load Collected data Trained model print('-------test-----------') class_number = { 'datasets/Market/datasets/pytorch/': 751, 'datasets/Duke/datasets/pytorch/': 702, 'datasets/CUHK03_detected/datasets/pytorch/': 767, 'datasets/CUHK03_labled/datasets/pytorch/': 767 } class_num = class_number[data_dir] part = 3 if not opt.mhn: model_structure = PCB(class_num, part) else: parts = opt.parts model_structure = MHN_smallPCB(class_num, parts, part) model_structure.model.fc = nn.Sequential() model_structure.model.avgpool = nn.Sequential() #print(model_structure) model = load_network(model_structure) ## load the learned params if not opt.mhn: model.model.avgpool = nn.Sequential() model.model.fc = nn.Sequential() # Remove the final fc layer and classifier layer model.classifier0.classifier = nn.Sequential() model.classifier1.classifier = nn.Sequential() model.classifier2.classifier = nn.Sequential()
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-----------') if opt.use_dense: model_structure = ft_net_dense(28310) else: model_structure = ft_net(28310) if opt.PCB: model_structure = PCB(28310) model_structure = nn.DataParallel(model_structure) model = load_network(model_structure) # Remove the final fc layer and classifier layer #print(model) if not opt.PCB: model.module.model.fc = nn.Sequential() # model.classifier = nn.Sequential() else: model = PCB_test(model) print(model) # Change to test mode
###################################################################### # Start training # --------------------------- # Record the training information # f = open(args.save_dir + 'train_log.txt', 'w') copyfile('./train.py', args.save_dir + '/train.py') copyfile('./model.py', args.save_dir + '/model.py') # step1: PCB training # stage = 'PCB' model = PCB(len(class_names)) if use_gpu: model = model.cuda() if is_parallel_train: import pdb #pdb.set_trace() model = nn.DataParallel(model, device_ids=gpu_ids) criterion = nn.CrossEntropyLoss() model = pcb_train(model, criterion, f, stage, 100) ############################ # step2&3: RPP training # if args.RPP:
print() return model # Save modeL def save_network(network, epoch_label): save_filename = 'net_%s.pth' % epoch_label save_path = os.path.join(opt.outf, save_filename) torch.save(network.cpu().state_dict(), save_path) if torch.cuda.is_available: network.cuda() # setting of train phase model = PCB(class_num=4768) # load pretrained para without classifier if use_gpu: model = model.cuda() # set the criterion triplet_selector = SemihardNegativeTripletSelector(opt.margin) criterion_tri = OnlineTripletLoss(opt.margin, triplet_selector) criterion_part = nn.CrossEntropyLoss() # criterion_part=CrossEntropyLabelSmooth(4768) criterion_center = CenterLoss(4768) criterion_focal = FocalLoss(gamma=2) # updating rule for parameter
for i in range(len(dataset_list)): dataset_path.append(image_datasets[dataset_list[i]].imgs) dataset_cam = [] dataset_label = [] dataset_filename = [] for i in range(len(dataset_list)): cam, label, filename = get_id(dataset_path[i]) dataset_cam.append(cam) dataset_label.append(label) dataset_filename.append(filename) ###################################################################### # Load Collected data Trained model print('-------test-----------') class_num = len(os.listdir(os.path.join(data_dir, 'train_all'))) model = PCB(class_num) if 'st' in opt.which_epoch: model = load_whole_network(model, name, opt.which_epoch + '_' + str(opt.net_loss_model)) else: model = load_whole_network(model, name, opt.which_epoch) model = model.eval() if use_gpu: model = model.cuda() # Extract feature dataset_feature = [] with torch.no_grad(): for i in range(len(dataset_list)): dataset_feature.append(extract_feature(model, dataloaders[dataset_list[i]])) result = {'gallery_f': dataset_feature[0].numpy(), 'gallery_label': dataset_label[0],
###################################################################### # Finetuning the convnet # ---------------------- # # Load a pretrainied model and reset final fully connected layer. # if opt.use_dense: model = ft_net_dense(len(class_names)) else: model = ft_net(len(class_names)) if opt.PCB: model = PCB(len(class_names)) print(model) if use_gpu: model = model.cuda() criterion = nn.CrossEntropyLoss() if not opt.PCB: ignored_params = list(map(id, model.model.fc.parameters() )) + list(map(id, model.classifier.parameters() )) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer_ft = optim.SGD([ {'params': base_params, 'lr': 0.01}, {'params': model.model.fc.parameters(), 'lr': 0.1}, {'params': model.classifier.parameters(), 'lr': 0.1}
###################################################################### # Finetuning the convnet # ---------------------- # # Load a pretrainied model and reset final fully connected layer. # if opt.use_dense: model = ft_net_dense(len(class_names)) else: model = ft_net(len(class_names)) if opt.PCB: model = PCB(len(class_names)) model_verif = verif_net() # print(model) # print(model_verif) if use_gpu: model = model.cuda() model_verif = model_verif.cuda() criterion = nn.CrossEntropyLoss() if not opt.PCB: ignored_params = list(map(id, model.model.fc.parameters())) + list(map(id, model.classifier.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer_ft = optim.SGD([
###################################################################### # Finetuning the convnet # ---------------------- # # Load a pretrainied model and reset final fully connected layer. # if opt.use_dense: model = ft_net_dense(len(class_names), opt.droprate) elif opt.use_NAS: model = ft_net_NAS(len(class_names), opt.droprate) else: model = ft_net(len(class_names), opt.droprate, opt.stride) if opt.PCB: model = PCB(len(class_names)) opt.nclasses = len(class_names) print(model) if not opt.PCB: ignored_params = list(map(id, model.classifier.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer_ft = optim.SGD([{ 'params': base_params, 'lr': 0.1 * opt.lr }, { 'params': model.classifier.parameters(), 'lr': opt.lr
else: x_embedded = x for i, t in enumerate(set(y)): idx = y == t plt.scatter(x_embedded[idx, 0], x_embedded[idx, 1], label=t) plt.title(title, fontsize=16) plt.legend() if save is not None: plt.savefig(save) with open(CONFIG_FILE, 'r') as stream: config = yaml.load(stream) nclasses = config['nclasses'] model = PCB(nclasses) model.load_state_dict(torch.load(WEIGHT_FILE)) model = model.eval().cuda() features = [] labels = [] fol_list = os.listdir(SAMPLE_FOLDER) for folder in fol_list: if folder == 'ALL': continue fol = os.path.join(SAMPLE_FOLDER, folder) if not os.path.isdir(fol): continue
def test(): print('-------test-----------') time_start = time.time() if opt.use_dense: model_structure = ft_net_dense(opt.nclasses) #elif opt.use_NAS: # model_structure = ft_net_NAS(opt.nclasses) else: model_structure = ft_net(opt.nclasses, stride=opt.stride) if opt.PCB: model_structure = PCB(opt.nclasses) #if opt.fp16: # model_structure = network_to_half(model_structure) model = load_network(model_structure) # Remove the final fc layer and classifier layer if opt.PCB: #if opt.fp16: # model = PCB_test(model[1]) #else: model = PCB_test(model) else: #if opt.fp16: #model[1].model.fc = nn.Sequential() #model[1].classifier = nn.Sequential() #else: model.classifier.classifier = nn.Sequential() # Change to test mode model = model.eval() if use_gpu: model = model.cuda() # Extract feature with torch.no_grad(): gallery_feature = extract_feature(model, dataloaders['gallery']) query_feature = extract_feature(model, dataloaders['query']) if opt.multi: mquery_feature = extract_feature(model, dataloaders['multi-query']) # Save to Matlab for check result = { 'gallery_f': gallery_feature.numpy(), 'gallery_label': gallery_label, 'gallery_cam': gallery_cam, 'query_f': query_feature.numpy(), 'query_label': query_label, 'query_cam': query_cam } scipy.io.savemat('pytorch_result.mat', result) print(opt.name) os.system('python evaluate_gpu.py') if opt.multi: result = { 'mquery_f': mquery_feature.numpy(), 'mquery_label': mquery_label, 'mquery_cam': mquery_cam } scipy.io.savemat('multi_query.mat', result) time_end = time.time() print('totally cost', time_end - time_start)
query_path = image_datasets['query'].imgs gallery_cam, gallery_label, gallery_frames = get_id(gallery_path) query_cam, query_label, query_frames = get_id(query_path) ###################################################################### # Load Collected data Trained model class_num = 751 print('-------test-----------') if opt.use_dense: model_structure = ft_net_dense(class_num) else: model_structure = ft_net(class_num) if opt.PCB: model_structure = PCB(class_num) model = load_network(model_structure) # Remove the final fc layer and classifier layer if not opt.PCB: model.model.fc = nn.Sequential() model.classifier = nn.Sequential() else: model = PCB_test(model) # Change to test mode model = model.eval() if use_gpu: model = model.cuda()
def __init__(self, gpuID, model_path): #################################################### # Options # ------------------ self.gpu_ids = 1 self.which_epoch = '59' self.batch_size = 1 self.use_dense = False self.use_PCB = False self.model_path = model_path self.class_num = 751 self.score_threshold = 0.9 self.confidence_threshold = 0.6 #################################################### # Set gpu # ------------------ use_gpu = torch.cuda.is_available() if not use_gpu: print('can not user gpu') exit() torch.cuda.set_device(self.gpu_ids) #################################################### # Load model # ------------------ print('load model...') if self.use_dense: model_structure = ft_net_dense(self.class_num) else: model_structure = ft_net(self.class_num) if self.use_PCB: model_structure = PCB(self.class_num) model = utils.load_network(model_structure, self.model_path, self.which_epoch) # Remove the final fc layer and classifier layer if not self.use_PCB: model.model.fc = nn.Sequential() model.classifier = nn.Sequential() else: model = PCB_test(model) model = model.eval() #print(model) if use_gpu: model = model.cuda() self.model = model #################################################### # Set Transform # ------------------ print('set transform...') self.data_transforms = transforms.Compose([ transforms.Resize((256, 128), interpolation=3), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) if self.use_PCB: self.data_transforms = transforms.Compose([ transforms.Resize((384, 192), interpolation=3), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])
###################################################################### # Finetuning the convnet # ---------------------- # # Load a pretrainied model and reset final fully connected layer. # if opt.use_dense: model = ft_net_dense(len(class_names)) else: #model = ft_net(len(class_names)) model = ResidualAttentionModel(len(class_names)) if opt.PCB: model = PCB(len(class_names)) print(model) if use_gpu: model = model.cuda() criterion = nn.CrossEntropyLoss() lr = 0.001 optimizer_ft = torch.optim.Adam(model.parameters(), lr=lr) if not opt.PCB: print("Here") #ignored_params = list(map(id, model.fc.parameters() )) + list(map(id, model.classifier.parameters() )) #base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) #optimizer_ft = optim.SGD([ # {'params': base_params, 'lr': 0.01},
labels.append(int(label)) camera_id.append(int(camera[0])) 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'])
def train(opt): version = torch.__version__ fp16 = opt.fp16 data_dir = opt.data_dir name = opt.name str_ids = opt.gpu_ids.split(',') gpu_ids = [] for str_id in str_ids: gid = int(str_id) if gid >= 0: gpu_ids.append(gid) # set gpu ids if len(gpu_ids) > 0: torch.cuda.set_device(gpu_ids[0]) cudnn.benchmark = True ###################################################################### # Load Data # --------- # transform_train_list = [ # transforms.RandomResizedCrop(size=128, scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC) transforms.Resize((256, 128), interpolation=3), transforms.Pad(10), transforms.RandomCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] transform_val_list = [ transforms.Resize(size=(256, 128), interpolation=3), # Image.BICUBIC transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.PCB: transform_train_list = [ transforms.Resize((384, 192), interpolation=3), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] transform_val_list = [ transforms.Resize(size=(384, 192), interpolation=3), # Image.BICUBIC transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.erasing_p > 0: transform_train_list = transform_train_list + [ RandomErasing(probability=opt.erasing_p, mean=[0.0, 0.0, 0.0]) ] if opt.color_jitter: transform_train_list = [ transforms.ColorJitter( brightness=0.1, contrast=0.1, saturation=0.1, hue=0) ] + transform_train_list # print(transform_train_list) data_transforms = { 'train': transforms.Compose(transform_train_list), 'val': transforms.Compose(transform_val_list), } train_all = '' if opt.train_all: train_all = '_all' image_datasets = {} image_datasets['train'] = datasets.ImageFolder( os.path.join(data_dir, 'train' + train_all), data_transforms['train']) image_datasets['val'] = datasets.ImageFolder(os.path.join(data_dir, 'val'), data_transforms['val']) dataloaders = { x: torch.utils.data.DataLoader(image_datasets[x], batch_size=opt.batchsize, shuffle=True, num_workers=8, pin_memory=True) # 8 workers may work faster for x in ['train', 'val'] } dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} class_names = image_datasets['train'].classes use_gpu = torch.cuda.is_available() #since = time.time() #inputs, classes = next(iter(dataloaders['train'])) #print('time used for loading data: %ds' %(time.time() - since)) ###################################################################### # Training the model # ------------------ # # Now, let's write a general function to train a model. Here, we will # illustrate: # # - Scheduling the learning rate # - Saving the best model # # In the following, parameter ``scheduler`` is an LR scheduler object from # ``torch.optim.lr_scheduler``. y_loss = {} # loss history y_loss['train'] = [] y_loss['val'] = [] y_err = {} y_err['train'] = [] y_err['val'] = [] def train_model(model, criterion, optimizer, scheduler, num_epochs=25): since = time.time() results = [] for epoch in range(num_epochs): print('Epoch {}/{}'.format(epoch, num_epochs - 1)) # Each epoch has a training and validation phase for phase in ['train', 'val']: if phase == 'train': scheduler.step() model.train(True) # Set model to training mode else: model.train(False) # Set model to evaluate mode running_loss = 0.0 running_corrects = 0.0 # Iterate over data. pbar = tqdm(dataloaders[phase]) for inputs, labels in pbar: # get the inputs now_batch_size, c, h, w = inputs.shape if now_batch_size < opt.batchsize: # skip the last batch continue # print(inputs.shape) # wrap them in Variable if use_gpu: inputs = Variable(inputs.cuda().detach()) labels = Variable(labels.cuda().detach()) else: inputs, labels = Variable(inputs), Variable(labels) # if we use low precision, input also need to be fp16 # if fp16: # inputs = inputs.half() # zero the parameter gradients optimizer.zero_grad() # forward if phase == 'val': with torch.no_grad(): outputs = model(inputs) else: outputs = model(inputs) if not opt.PCB: _, preds = torch.max(outputs.data, 1) loss = criterion(outputs, labels) else: part = {} sm = nn.Softmax(dim=1) num_part = 6 for i in range(num_part): part[i] = outputs[i] score = sm(part[0]) + sm(part[1]) + sm(part[2]) + sm( part[3]) + sm(part[4]) + sm(part[5]) _, preds = torch.max(score.data, 1) loss = criterion(part[0], labels) for i in range(num_part - 1): loss += criterion(part[i + 1], labels) # backward + optimize only if in training phase if phase == 'train': if fp16: # we use optimier to backward loss with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() optimizer.step() # statistics if int(version[0]) > 0 or int( version[2] ) > 3: # for the new version like 0.4.0, 0.5.0 and 1.0.0 running_loss += loss.item() * now_batch_size else: # for the old version like 0.3.0 and 0.3.1 running_loss += loss.data[0] * now_batch_size running_corrects += float(torch.sum(preds == labels.data)) pbar.set_description( desc='loss: {:.4f}'.format(loss.item())) epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects / dataset_sizes[phase] print('\r\n{} Loss: {:.4f} Acc: {:.4f}'.format( phase, epoch_loss, epoch_acc)) logging.info('epoch: {}, {} Loss: {:.4f} Acc: {:.4f}'.format( epoch, phase, epoch_loss, epoch_acc)) y_loss[phase].append(epoch_loss) y_err[phase].append(1.0 - epoch_acc) # deep copy the model if phase == 'val': results.append({ 'epoch': epoch, 'trainLoss': y_loss['train'][-1], 'trainError': y_err['train'][-1], 'valLoss': y_loss['val'][-1], 'valError': y_err['val'][-1] }) last_model_wts = model.state_dict() if epoch % 10 == 9: save_network(model, epoch) draw_curve(epoch) write_to_csv(results) time_elapsed = time.time() - since print('\r\nTraining complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) print() time_elapsed = time.time() - since print('\r\nTraining complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) # print('Best val Acc: {:4f}'.format(best_acc)) # load best model weights model.load_state_dict(last_model_wts) save_network(model, 'last') return model ###################################################################### # Draw Curve # --------------------------- x_epoch = [] fig = plt.figure() ax0 = fig.add_subplot(121, title="loss") ax1 = fig.add_subplot(122, title="top1err") def draw_curve(current_epoch): x_epoch.append(current_epoch) ax0.plot(x_epoch, y_loss['train'], 'bo-', label='train') ax0.plot(x_epoch, y_loss['val'], 'ro-', label='val') ax1.plot(x_epoch, y_err['train'], 'bo-', label='train') ax1.plot(x_epoch, y_err['val'], 'ro-', label='val') if current_epoch == 0: ax0.legend() ax1.legend() fig.savefig(os.path.join('./model', name, 'train.jpg')) def write_to_csv(results): path = os.path.join('./model', name, 'result.csv') with open(path, 'w', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=list(results[0].keys())) writer.writeheader() writer.writerows(results) ###################################################################### # Save model # --------------------------- def save_network(network, epoch_label): save_filename = 'net_%s.pth' % epoch_label rpth = os.path.join('./model', name, 'Model Files') if not os.path.exists(rpth): os.makedirs(rpth) save_path = os.path.join(rpth, save_filename) torch.save(network.cpu().state_dict(), save_path) if torch.cuda.is_available(): network.cuda(gpu_ids[0]) ###################################################################### # Finetuning the convnet # ---------------------- # # Load a pretrainied model and reset final fully connected layer. # if opt.use_dense: model = ft_net_dense(len(class_names), opt.droprate) else: model = ft_net(len(class_names), opt.droprate, opt.stride) if opt.PCB: model = PCB(len(class_names)) opt.nclasses = len(class_names) print(model) print('model loaded') if not opt.PCB: ignored_params = list(map(id, model.model.fc.parameters())) + list( map(id, model.classifier.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer_ft = optim.SGD([{ 'params': base_params, 'lr': 0.1 * opt.lr }, { 'params': model.model.fc.parameters(), 'lr': opt.lr }, { 'params': model.classifier.parameters(), 'lr': opt.lr }], weight_decay=5e-4, momentum=0.9, nesterov=True) else: ignored_params = list(map(id, model.model.fc.parameters())) ignored_params += ( list(map(id, model.classifier0.parameters())) + list(map(id, model.classifier1.parameters())) + list(map(id, model.classifier2.parameters())) + list(map(id, model.classifier3.parameters())) + list(map(id, model.classifier4.parameters())) + list(map(id, model.classifier5.parameters())) # +list(map(id, model.classifier6.parameters() )) # +list(map(id, model.classifier7.parameters() )) ) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer_ft = optim.SGD( [ { 'params': base_params, 'lr': 0.1 * opt.lr }, { 'params': model.model.fc.parameters(), 'lr': opt.lr }, { 'params': model.classifier0.parameters(), 'lr': opt.lr }, { 'params': model.classifier1.parameters(), 'lr': opt.lr }, { 'params': model.classifier2.parameters(), 'lr': opt.lr }, { 'params': model.classifier3.parameters(), 'lr': opt.lr }, { 'params': model.classifier4.parameters(), 'lr': opt.lr }, { 'params': model.classifier5.parameters(), 'lr': opt.lr }, # {'params': model.classifier6.parameters(), 'lr': 0.01}, # {'params': model.classifier7.parameters(), 'lr': 0.01} ], weight_decay=5e-4, momentum=0.9, nesterov=True) # Decay LR by a factor of 0.1 every 40 epochs exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=40, gamma=0.1) ###################################################################### # Train and evaluate # ^^^^^^^^^^^^^^^^^^ # # It should take around 1-2 hours on GPU. # dir_name = os.path.join('./model', name) if not os.path.isdir(dir_name): os.mkdir(dir_name) # record every run copyfile('./train.py', dir_name + '/train.py') copyfile('./model.py', dir_name + '/model.py') # save opts with open('%s/opts.yaml' % dir_name, 'w') as fp: yaml.dump(vars(opt), fp, default_flow_style=False) # model to gpu model = model.cuda() if fp16: # model = network_to_half(model) # optimizer_ft = FP16_Optimizer(optimizer_ft, static_loss_scale = 128.0) model, optimizer_ft = amp.initialize(model, optimizer_ft, opt_level="O1") criterion = losses.DualLoss() model = train_model(model, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=60) # # if __name__ == "__main__": # train(opt)
if opt.multi: mquery_path = image_datasets['multi-query'].imgs mquery_cam, mquery_label = get_id(mquery_path) ###################################################################### # Load Collected data Trained model print('-------test-----------') if opt.use_dense: model_structure = ft_net_dense(opt.nclasses) elif opt.use_NAS: model_structure = ft_net_NAS(opt.nclasses) else: model_structure = ft_net(opt.nclasses, stride=opt.stride) if opt.PCB: model_structure = PCB(opt.nclasses) #if opt.fp16: # model_structure = network_to_half(model_structure) model = load_network(model_structure) # Remove the final fc layer and classifier layer if opt.PCB: #if opt.fp16: # model = PCB_test(model[1]) #else: model = PCB_test(model) else: #if opt.fp16: #model[1].model.fc = nn.Sequential()
save_whole_network(model, name, 'last' + '_' + str(opt.net_loss_model)) return model ###################################################################### # Train and evaluate # ^^^^^^^^^^^^^^^^^^ # # It should take around 1-2 hours on GPU. # dir_name = os.path.join('./model', name) if not os.path.exists('model'): os.mkdir('model') print('class_num = %d' % (class_num)) model = PCB(class_num, train=True) if use_gpu: model.cuda() # print('model structure') # print(model) criterion = nn.CrossEntropyLoss() criterion_soft = SoftLabelLoss() classifier_id = (list(map(id, model.classifier0.parameters())) + list(map(id, model.classifier1.parameters())) + list(map(id, model.classifier2.parameters())) + list(map(id, model.classifier3.parameters())) + list(map(id, model.classifier4.parameters())) + list(map(id, model.classifier5.parameters()))) classifier_params = filter(lambda p: id(p) in classifier_id,