def __init__(self, channels, blocks, linears, num_class, flatten=False, feat_dim=None, gpu_flag=True): super(CNN, self).__init__() self.channels = channels self.flatten = flatten if flatten: self.flatten_shape = calculateShape(self.channels) self.feat_dim = feat_dim else: self.flatten_shape = self.channels[-1] self.feat_dim = self.channels[-1] # self.linears = [self.flatten_shape] + linears + [num_class] self.linears = [self.feat_dim] + linears + [num_class] print(self.linears) self.num_class = num_class self.inchannels = 64 self.blocks = blocks self._norm_layer = nn.BatchNorm2d self.cnn = self._make_cnn(self.channels, self.blocks) self.classifier = self._make_classifier(self.linears) # self.layers = [] # self.layers.append(self._make_cnn(self.channels, self.blocks)) # self.layers.append(self._make_classifier(self.linears)) # embedding for center loss self.emb_layer = nn.Sequential( # nn.BatchNorm1d(self.linears[0]), nn.Linear(self.flatten_shape, self.feat_dim, bias=False), ) self.emb_avg = nn.AdaptiveAvgPool2d((1, 1)) self.closs_criterion = CenterLoss(self.num_class, self.feat_dim, use_gpu=gpu_flag) self.ce_criterion = nn.CrossEntropyLoss()
def __init__(self, input_dim, layers, feat_dim, num_class, gpu_flag=True): super(MLP, self).__init__() self.CE = torch.nn.CrossEntropyLoss() self.softmax = nn.Softmax(dim=1) self.input_dim = input_dim self.layers = [self.input_dim] + layers + [num_class] self.num_class = num_class self.feat_dim = feat_dim self.emb_linears = nn.Linear(self.layers[0], self.layers[1], bias=True) self.label_linears = [] for i in range(1, len(self.layers) - 1): self.label_linears.append(nn.BatchNorm1d(self.layers[i])) self.label_linears.append(nn.PReLU()) self.label_linears.append( nn.Linear(self.layers[i], self.layers[i + 1], bias=True)) #self.linears.append(nn.Linear(self.layers[-2], self.layers[-1], bias=True)) self.label_linears = nn.Sequential(*self.label_linears) self.closs_criterion = CenterLoss(self.num_class, self.feat_dim, use_gpu=gpu_flag) self.ce_criterion = nn.CrossEntropyLoss()
def train(parameters: Dict[str, float]) -> nn.Module: global args print("====", args.focus, "=====") torch.manual_seed(args.seed) # args.gpu_devices = "0,1" os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) dataset = data_manager.init_dataset(name=args.dataset, sampling= args.sampling) transform_test = transforms.Compose([ transforms.Resize((args.height, args.width), interpolation=3), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) pin_memory = True if use_gpu else False transform_train = transforms.Compose([ transforms.Resize((args.height, args.width), interpolation=3), transforms.RandomHorizontalFlip(p=0.5), transforms.Pad(10), Random2DTranslation(args.height, args.width), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) batch_size = int(round(parameters.get("batch_size", 32) )) base_learning_rate = 0.00035 # weight_decay = 0.0005 alpha = parameters.get("alpha", 1.2) sigma = parameters.get("sigma", 0.8) l = parameters.get("l", 0.5) beta_ratio = parameters.get("beta_ratio", 0.5) gamma = parameters.get("gamma", 0.1) margin = parameters.get("margin", 0.3) weight_decay = parameters.get("weight_decay", 0.0005) lamb = 0.3 num_instances = 4 pin_memory = True trainloader = DataLoader( VideoDataset(dataset.train, seq_len=args.seq_len, sample='random',transform=transform_train), sampler=RandomIdentitySampler(dataset.train, num_instances=args.num_instances), batch_size=batch_size, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) if args.dataset == 'mars_subset' : validation_loader = DataLoader( VideoDataset(dataset.val, seq_len=8, sample='random', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) else: queryloader = DataLoader( VideoDataset(dataset.val_query, seq_len=args.seq_len, sample='dense_subset', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( VideoDataset(dataset.val_gallery, seq_len=args.seq_len, sample='dense_subset', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) criterion_htri = TripletLoss(margin, 'cosine') criterion_xent = CrossEntropyLabelSmooth(dataset.num_train_pids) criterion_center_loss = CenterLoss(use_gpu=1) criterion_osm_caa = OSM_CAA_Loss(alpha=alpha , l=l , osm_sigma=sigma ) args.arch = "ResNet50ta_bt" model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) if use_gpu: model = nn.DataParallel(model).cuda() params = [] for key, value in model.named_parameters(): if not value.requires_grad: continue lr = base_learning_rate weight_decay = weight_decay params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}] optimizer = torch.optim.Adam(params) scheduler = WarmupMultiStepLR(optimizer, milestones=[40, 70], gamma=gamma, warmup_factor=0.01, warmup_iters=10) optimizer_center = torch.optim.SGD(criterion_center_loss.parameters(), lr=0.5) start_epoch = args.start_epoch best_rank1 = -np.inf num_epochs = 121 if 'mars' not in args.dataset : num_epochs = 121 # test_rerank(model, queryloader, galleryloader, args.pool, use_gpu, lamb=lamb , parameters=parameters) for epoch in range (num_epochs): vals = train_model(model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu , optimizer_center , criterion_center_loss, criterion_osm_caa, beta_ratio) if math.isnan(vals[0]): return 0 scheduler.step() if epoch % 40 ==0 : print("TripletLoss {:.6f} OSM Loss {:.6f} Cross_entropy {:.6f} Total Loss {:.6f} ".format(vals[1] , vals[3] , vals[1] , vals[0])) if args.dataset == 'mars_subset' : result1 = test_validation(model, validation_loader, args.pool, use_gpu, parameters=parameters) del validation_loader else: result1= test_rerank(model, queryloader, galleryloader, args.pool, use_gpu, lamb=lamb , parameters=parameters) del queryloader del galleryloader del trainloader del model del criterion_htri del criterion_xent del criterion_center_loss del criterion_osm_caa del optimizer del optimizer_center del scheduler return result1
def main(): #GENERAL torch.cuda.empty_cache() root = "/home/kuru/Desktop/veri-gms-master/" train_dir = '/home/kuru/Desktop/veri-gms-master/VeRispan/image_train/' source = {'verispan'} target = {'verispan'} workers = 4 height = 320 width = 320 train_sampler = 'RandomSampler' #AUGMENTATION random_erase = True jitter = True aug = True #OPTIMIZATION opt = 'adam' lr = 0.001 weight_decay = 5e-4 momentum = 0.9 sgd_damp = 0.0 nesterov = True warmup_factor = 0.01 warmup_method = 'linear' STEPS = (30, 60) GAMMA = 0.1 WARMUP_FACTOR = 0.01 WARMUP_EPOCHS = 10 WARMUP_METHOD = 'linear' #HYPERPARAMETER max_epoch = 80 start = 0 train_batch_size = 16 test_batch_size = 50 #SCHEDULER lr_scheduler = 'multi_step' stepsize = [30, 60] gamma = 0.1 #LOSS margin = 0.3 num_instances = 4 lambda_tri = 1 #MODEL #arch = 'resnet101' arch = 'resnet50_ibn_a' no_pretrained = False #TEST SETTINGS #load_weights = '/home/kuru/Desktop/veri-gms-master/IBN-Net_pytorch0.4.1/resnet101_ibn_a.pth' #load_weights = '/home/kuru/Desktop/veri-gms-master/IBN-Net_pytorch0.4.1/resnet101_ibn_a.pth' load_weights = '/home/kuru/Desktop/veri-gms-master/IBN-Net_pytorch0.4.1/resnet50_ibn_a.pth' #load_weights = None start_eval = 0 eval_freq = -1 num_classes = 776 feat_dim = 2048 CENTER_LR = 0.5 CENTER_LOSS_WEIGHT = 0.0005 center_criterion = CenterLoss(num_classes=num_classes, feat_dim=feat_dim, use_gpu=True) optimizer_center = torch.optim.SGD(center_criterion.parameters(), lr=CENTER_LR) #MISC use_gpu = True #use_gpu = False print_freq = 10 seed = 1 resume = '' save_dir = '/home/kuru/Desktop/veri-gms-master_noise/spanningtree_veri_pure/' gpu_id = 0, 1 vis_rank = True query_remove = True evaluate = False dataset_kwargs = { 'source_names': source, 'target_names': target, 'root': root, 'height': height, 'width': width, 'train_batch_size': train_batch_size, 'test_batch_size': test_batch_size, 'train_sampler': train_sampler, 'random_erase': random_erase, 'color_jitter': jitter, 'color_aug': aug } transform_kwargs = { 'height': height, 'width': width, 'random_erase': random_erase, 'color_jitter': jitter, 'color_aug': aug } optimizer_kwargs = { 'optim': opt, 'lr': lr, 'weight_decay': weight_decay, 'momentum': momentum, 'sgd_dampening': sgd_damp, 'sgd_nesterov': nesterov } lr_scheduler_kwargs = { 'lr_scheduler': lr_scheduler, 'stepsize': stepsize, 'gamma': gamma } use_gpu = torch.cuda.is_available() log_name = 'log_test.txt' if evaluate else 'log_train.txt' sys.stdout = Logger(osp.join(save_dir, log_name)) print('Currently using GPU ', gpu_id) cudnn.benchmark = True print('Initializing image data manager') #dataset = init_imgreid_dataset(root='/home/kuru/Desktop/veri-gms-master/', name='veri') dataset = init_imgreid_dataset(root='/home/kuru/Desktop/veri-gms-master/', name='verispan') train = [] num_train_pids = 0 num_train_cams = 0 print(len(dataset.train)) for img_path, pid, camid, subid, countid in dataset.train: #print(img_path) path = img_path[56:90 + 6] #print(path) folder = path[1:4] #print(folder) #print(img_path, pid, camid,subid,countid) pid += num_train_pids camid += num_train_cams newidd = 0 train.append((path, folder, pid, camid, subid, countid)) #print(train) #break num_train_pids += dataset.num_train_pids num_train_cams += dataset.num_train_cams pid = 0 pidx = {} for img_path, pid, camid, subid, countid in dataset.train: path = img_path[56:90 + 6] folder = path[1:4] pidx[folder] = pid pid += 1 #print(pidx) sub = [] final = 0 xx = dataset.train newids = [] print(train[0:2]) train2 = {} for k in range(0, 770): for img_path, pid, camid, subid, countid in dataset.train: if k == pid: newid = final + subid sub.append(newid) #print(pid,subid,newid) newids.append(newid) train2[img_path] = newid #print(img_path, pid, camid, subid, countid, newid) final = max(sub) #print(final) print(len(newids), final) #train=train2 #print(train2) train3 = [] for img_path, pid, camid, subid, countid in dataset.train: #print(img_path,pid,train2[img_path]) path = img_path[56:90 + 6] #print(path) folder = path[1:4] newid = train2[img_path] #print((path, folder, pid, camid, subid, countid,newid )) train3.append((path, folder, pid, camid, subid, countid, newid)) train = train3 # for (path, folder, pid, camid, subid, countid,newid) in train: # print(path, folder) #path = '/home/kuru/Desktop/adhi/veri-final-draft-master_noise/gmsNoise776/' path = '/home/kuru/Desktop/veri-gms-master/gms/' pkl = {} #pkl[0] = pickle.load('/home/kuru/Desktop/veri-gms-master/gms/620.pkl') entries = os.listdir(path) for name in entries: f = open((path + name), 'rb') ccc = (path + name) #print(ccc) if name == 'featureMatrix.pkl': s = name[0:13] else: s = name[0:3] #print(s) #with open (ccc,"rb") as ff: # pkl[s] = pickle.load(ff) #print(pkl[s]) pkl[s] = pickle.load(f) f.close #print(len(pkl)) print('=> pickle indexing') data_index = search(pkl) print(len(data_index)) transform_t = train_transforms(**transform_kwargs) #print(train[0],train[10]) #data_tfr = vd(pkl_file='index.pkl', dataset = train, root_dir='/home/kuru/Desktop/veri-gms-master/VeRi/image_train/', transform=transform_t) data_tfr = vdspan( pkl_file='index_veryspan.pkl', dataset=train, root_dir='/home/kuru/Desktop/veri-gms-master/VeRispan/image_train/', transform=transform_t) #print(data_tfr) #print(trainloader) #data_tfr2=list(data_tfr) print("lllllllllllllllllllllllllllllllllllllllllllline 433") df2 = [] data_tfr_old = data_tfr for (img, label, index, pid, cid, subid, countid, newid) in data_tfr: #print((img,label,index,pid, cid,subid,countid,newid) ) #print("datframe",(label)) #print(countid) if countid > 4: #print(countid) df2.append((img, label, index, pid, cid, subid, countid, newid)) print("filtered final trainset length", len(df2)) data_tfr = df2 # with open('df2noise_ex.pkl', 'wb') as handle: # b = pickle.dump(df2, handle, protocol=pickle.HIGHEST_PROTOCOL) # with open('df2noise.pkl', 'rb') as handle: # df2 = pickle.load(handle) # data_tfr=df2 # for (img,label,index,pid, cid,subid,countid,newid) in data_tfr : # print("datframe",(label)) #data_tfr = vdspansort( dataset = train, root_dir='/home/kuru/Desktop/veri-gms-master_noise/VeRispan/image_train/', transform=transform_t) #trainloader = DataLoader(df2, sampler=None,batch_size=train_batch_size, shuffle=True, num_workers=workers,pin_memory=True, drop_last=True) trainloader = DataLoader(data_tfr, sampler=None, batch_size=train_batch_size, shuffle=True, num_workers=workers, pin_memory=True, drop_last=True) for batch_idx, (img, label, index, pid, cid, subid, countid, newid) in enumerate(trainloader): #print("trainloader",batch_idx, (label,index,pid, cid,subid,countid,newid)) print("trainloader", batch_idx, (label)) break print('Initializing test data manager') dm = ImageDataManager(use_gpu, **dataset_kwargs) testloader_dict = dm.return_dataloaders() print('Initializing model: {}'.format(arch)) model = models.init_model(name=arch, num_classes=num_train_pids, loss={'xent', 'htri'}, pretrained=not no_pretrained, last_stride=2) print('Model size: {:.3f} M'.format(count_num_param(model))) if load_weights is not None: print("weights loaded") load_pretrained_weights(model, load_weights) #checkpoint = torch.load('/home/kuru/Desktop/veri-gms-master/logg/model.pth.tar-19') #model._load_from_state_dict(checkpoint['state_dict']) #model.load_state_dict(checkpoint['state_dict']) #optimizer.load_state_dict(checkpoint['optimizer']) #print(checkpoint['epoch']) #print(checkpoint['rank1']) os.environ['CUDA_VISIBLE_DEVICES'] = '0' print(torch.cuda.device_count()) model = nn.DataParallel(model).cuda() if use_gpu else model optimizer = init_optimizer(model, **optimizer_kwargs) #optimizer = init_optimizer(model) #optimizer.load_state_dict(checkpoint['optimizer']) scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs) # scheduler = WarmupMultiStepLR(optimizer, STEPS, GAMMA, # WARMUP_FACTOR, # WARMUP_EPOCHS, WARMUP_METHOD) criterion_xent = CrossEntropyLoss(num_classes=num_train_pids, use_gpu=use_gpu, label_smooth=True) criterion_htri = TripletLoss(margin=margin) ranking_loss = nn.MarginRankingLoss(margin=margin) if evaluate: print('Evaluate only') for name in target: print('Evaluating {} ...'.format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] _, distmat = test(model, queryloader, galleryloader, train_batch_size, use_gpu, return_distmat=True) if vis_rank: visualize_ranked_results(distmat, dm.return_testdataset_by_name(name), save_dir=osp.join( save_dir, 'ranked_results', name), topk=20) return time_start = time.time() ranklogger = RankLogger(source, target) # # checkpoint = torch.load('/home/kuru/Desktop/market_all/ibna_model/model.pth.tar-79') # # model.load_state_dict(checkpoint['state_dict']) # # optimizer.load_state_dict(checkpoint['optimizer']) # # print(checkpoint['epoch']) # # start_epoch=checkpoint['epoch'] # # start=start_epoch # checkpoint = torch.load('/home/kuru/Desktop/veri-gms-master/spanningtreeveri/model.pth.tar-2') # model.load_state_dict(checkpoint['state_dict']) # optimizer.load_state_dict(checkpoint['optimizer']) # print(checkpoint['epoch']) # start_epoch=checkpoint['epoch'] # start=start_epoch ##start_epoch=resume_from_checkpoint('/home/kuru/Desktop/veri-gms-master/logg/model.pth.tar-20', model, optimizer=None) print('=> Start training') for epoch in range(start, max_epoch): print(epoch, scheduler.get_lr()[0]) #print( torch.cuda.memory_allocated(0)) losses = AverageMeter() #xent_losses = AverageMeter() htri_losses = AverageMeter() accs = AverageMeter() batch_time = AverageMeter() xent_losses = AverageMeter() model.train() for p in model.parameters(): p.requires_grad = True # open all layers end = time.time() for batch_idx, (img, label, index, pid, cid, subid, countid, newid) in enumerate(trainloader): trainX, trainY = torch.zeros( (train_batch_size * 3, 3, height, width), dtype=torch.float32), torch.zeros((train_batch_size * 3), dtype=torch.int64) #pids = torch.zeros((batch_size*3), dtype = torch.int16) #batchcount=0 for i in range(train_batch_size): if (countid[i] > 4): #batchcount=batchcount+1 #print("dfdsfs") labelx = label[i] indexx = index[i] cidx = pid[i] if indexx > len(pkl[labelx]) - 1: indexx = len(pkl[labelx]) - 1 #maxx = np.argmax(pkl[labelx][indexx]) a = pkl[labelx][indexx] minpos = np.argmin(ma.masked_where(a == 0, a)) # print(len(a)) # print(a) # print(ma.masked_where(a==0, a)) # print(labelx,index,pid,cidx,minpos) # print(np.array(data_index).shape) # print(data_index[cidx][1]) pos_dic = data_tfr_old[data_index[cidx][1] + minpos] #print('posdic', pos_dic) neg_label = int(labelx) while True: neg_label = random.choice(range(1, 770)) #print(neg_label) if neg_label is not int(labelx) and os.path.isdir( os.path.join( '/home/kuru/Desktop/veri-gms-master_noise/veriNoise_train_spanning_folder', strint(neg_label))) is True: break negative_label = strint(neg_label) neg_cid = pidx[negative_label] neg_index = random.choice( range(0, len(pkl[negative_label]))) #print(negative_label,neg_cid,neg_index,data_index[neg_cid] ) neg_dic = data_tfr_old[data_index[neg_cid][1] + neg_index] #print('negdic', neg_dic) trainX[i] = img[i] trainX[i + train_batch_size] = pos_dic[0] trainX[i + (train_batch_size * 2)] = neg_dic[0] trainY[i] = cidx trainY[i + train_batch_size] = pos_dic[3] trainY[i + (train_batch_size * 2)] = neg_dic[3] # trainY[i+train_batch_size] = pos_dic[7] # trainY[i+(train_batch_size*2)] = neg_dic[7] #break # else: # print("skiped",countid[i],subid[i],label[i]) #break #print(batchcount) trainX = trainX.cuda() trainY = trainY.cuda() outputs, features = model(trainX) xent_loss = criterion_xent(outputs[0:train_batch_size], trainY[0:train_batch_size]) htri_loss = criterion_htri(features, trainY) centerloss = CENTER_LOSS_WEIGHT * center_criterion( features, trainY) #tri_loss = ranking_loss(features) #ent_loss = xent_loss(outputs[0:batch_size], trainY[0:batch_size], num_train_pids) loss = htri_loss + xent_loss + centerloss loss = htri_loss + xent_loss optimizer.zero_grad() optimizer_center.zero_grad() loss.backward() optimizer.step() # for param in center_criterion.parameters(): # param.grad.data *= (1. /CENTER_LOSS_WEIGHT) # optimizer_center.step() for param_group in optimizer.param_groups: #print(param_group['lr'] ) lrrr = str(param_group['lr']) batch_time.update(time.time() - end) losses.update(loss.item(), trainY.size(0)) htri_losses.update(htri_loss.item(), trainY.size(0)) xent_losses.update(xent_loss.item(), trainY.size(0)) accs.update( accuracy(outputs[0:train_batch_size], trainY[0:train_batch_size])[0]) if (batch_idx) % 50 == 0: print('Train ', end=" ") print('Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' 'TriLoss {loss.val:.4f} ({loss.avg:.4f})\t' 'XLoss {xloss.val:.4f} ({xloss.avg:.4f})\t' 'OveralLoss {oloss.val:.4f} ({oloss.avg:.4f})\t' 'Acc {acc.val:.2f} ({acc.avg:.2f})\t' 'lr {lrrr} \t'.format( epoch + 1, batch_idx + 1, len(trainloader), batch_time=batch_time, loss=htri_losses, xloss=xent_losses, oloss=losses, acc=accs, lrrr=lrrr, )) end = time.time() # del loss # del htri_loss # del xent_loss # del htri_losses # del losses # del outputs # del features # del accs # del trainX # del trainY scheduler.step() print('=> Test') save_checkpoint( { 'state_dict': model.state_dict(), #'rank1': rank1, 'epoch': epoch + 1, 'arch': arch, 'optimizer': optimizer.state_dict(), }, save_dir) GPUtil.showUtilization() print(torch.cuda.memory_allocated(), torch.cuda.memory_cached()) for name in target: print('Evaluating {} ...'.format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] rank1, distmat = test(model, queryloader, galleryloader, test_batch_size, use_gpu) ranklogger.write(name, epoch + 1, rank1) rank2, distmat2 = test_rerank(model, queryloader, galleryloader, test_batch_size, use_gpu) ranklogger.write(name, epoch + 1, rank2) del queryloader del galleryloader del distmat print(torch.cuda.memory_allocated(), torch.cuda.memory_cached()) torch.cuda.empty_cache() if (epoch + 1) == max_epoch: #if (epoch + 1) % 10 == 0: print('=> Test') save_checkpoint( { 'state_dict': model.state_dict(), 'rank1': rank1, 'epoch': epoch + 1, 'arch': arch, 'optimizer': optimizer.state_dict(), }, save_dir) for name in target: print('Evaluating {} ...'.format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] rank1, distmat = test(model, queryloader, galleryloader, test_batch_size, use_gpu) ranklogger.write(name, epoch + 1, rank1) # del queryloader # del galleryloader # del distmat if vis_rank: visualize_ranked_results( distmat, dm.return_testdataset_by_name(name), save_dir=osp.join(save_dir, 'ranked_results', name), topk=20)
l = 0.5241677630566622 margin = 0.040520629258433416 beta_ratio = 0.7103921571238655 gamma = 0.368667605025003 weight_decay = 0.014055481861393148 args.arch = "ResNet50ta_bt" model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) print("Model size: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) criterion_htri = TripletLoss(margin, 'cosine') criterion_xent = CrossEntropyLabelSmooth(dataset.num_train_pids) criterion_center_loss = CenterLoss(use_gpu=use_gpu) if args.use_OSMCAA: print("USING OSM LOSS") print("config, alpha = %f sigma = %f l=%f" % (alpha, sigma, l)) criterion_osm_caa = OSM_CAA_Loss(alpha=alpha, l=l, osm_sigma=sigma) else: criterion_osm_caa = None if args.cl_centers: print("USING CL CENTERS") print("config, alpha = %f sigma = %f l=%f" % (alpha, sigma, l)) criterion_osm_caa = OSM_CAA_Loss(alpha=alpha, l=l, osm_sigma=sigma) base_learning_rate = 0.00035 params = []
def train_save(model, trainset_loader, testset_loader, opt, epoch=5, loss_cri='Focal', save_interval=4000, log_interval=100, device='cpu', save_ep=False): os.makedirs('./models/', exist_ok=True) os.makedirs('./models/{}/'.format(opt.model_name), exist_ok=True) if opt.optim == 'Adam': optimizer = optim.Adam(model.parameters(), lr=opt.lr, betas=(0.9, 0.999)) elif opt.optim == 'SGD': optimizer = optim.SGD(model.parameters(), lr=opt.lr, momentum=0.9) if loss_cri == 'Focal': criterion = FocalLoss(alpha=opt.alpha, gamma=opt.gamma) elif loss_cri == 'Triplet': criterion = TripletLoss(margin=opt.margin) elif loss_cri == 'Center': criterion = CenterLoss() # criterion = nn.CrossEntropyLoss() iteration = 0 for ep in range(epoch): model.train() # set training mode for batch_idx, (features, labels) in enumerate(trainset_loader): features, labels = features.to(device), labels.squeeze(1).to( device) optimizer.zero_grad() output = model(features) loss = criterion(output, labels) loss.backward() optimizer.step() if iteration % log_interval == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( ep, batch_idx * len(features), len(trainset_loader.dataset), 100. * batch_idx / len(trainset_loader), loss.item())) if iteration % save_interval == 0 and iteration > 0: save_checkpoint( './models/{}/{}_backup.pth'.format(opt.model_name, opt.model_name), model, optimizer) iteration += 1 if opt.save_ep: save_checkpoint( './models/{}/{}_epoch_{}.pth'.format(opt.model_name, opt.model_name, ep), model, optimizer) eval(model, testset_loader, opt, criterion, threshold=opt.threshold, threshold_2=opt.threshold_2, epoch=ep) # save the final model save_checkpoint( './models/{}/{}_final.pth'.format(opt.model_name, opt.model_name), model, optimizer)
def _get_losses(self): return CenterLoss()
def train(): ## setup torch.multiprocessing.set_sharing_strategy('file_system') if not os.path.exists('./res'): os.makedirs('./res') ## model and loss logger.info('setting up backbone model and loss') net = EmbedNetwork().cuda() net = nn.DataParallel(net) print(net) triplet_loss = TripletLoss( margin=0.3).cuda() # no margin means soft-margin BNNeck = ClassBlock(2048, 1501).cuda() BNNeck = nn.DataParallel(BNNeck) ## optimizer logger.info('creating optimizer') optim = AdamOptimWrapper(net.parameters(), lr=3e-4, wd=0, t0=15000, t1=25000) ## dataloader selector = BatchHardTripletSelector() ds = Market1501( 'E://graduation thesis//triplet-reid-pytorch-master//triplet-reid-pytorch-master//datasets//Market-1501-v15.09.15//Market-1501-v15.09.15//bounding_box_train', is_train=True) sampler = BatchSampler(ds, 9, 4) dl = DataLoader(ds, batch_sampler=sampler, num_workers=0) diter = iter(dl) ## train logger.info('start training ...') loss_avg = [] loss1_avg = [] loss2_avg = [] loss3_avg = [] count = 0 t_start = time.time() while True: try: imgs, lbs, _ = next(diter) except StopIteration: diter = iter(dl) imgs, lbs, _ = next(diter) #criterion = nn.CrossEntropyLoss().cuda() criterion = CrossEntropyLabelSmooth(num_classes=1501) center_criterion = CenterLoss(num_classes=1051, feat_dim=2048, use_gpu=True) net.train() imgs = imgs.cuda() lbs = lbs.cuda() # for name in net.state_dict(): # print("net parameters:", name) optim.zero_grad() embds = net(imgs) anchor, positives, negatives = selector(embds, lbs) BNNeck.train() # for name in BNNeck.state_dict(): # print("BNNeck parameters:", name) #print(BNNeck) classifier = [] classifier += [nn.Linear(2048, 1501)] classifier = nn.Sequential(*classifier) classifier.apply(weights_init_classifier) classifier = classifier.cuda() x = torch.squeeze(embds) BNNeck1 = BNNeck(x) #classifier = torch.autograd.Variable(classifier) classifier = classifier(BNNeck1) loss1 = triplet_loss(anchor, positives, negatives) loss2 = criterion(classifier, lbs) loss3 = center_criterion(embds, lbs) loss = loss1 + loss2 + 0.0005 * loss3 loss.backward() optim.step() loss_avg.append(loss.detach().cpu().numpy()) loss1_avg.append(loss1.detach().cpu().numpy()) loss2_avg.append(loss2.detach().cpu().numpy()) loss3_avg.append(loss3.detach().cpu().numpy()) #loss1_avg.append(loss1.detach().cpu().numpy()) if count % 20 == 0 and count != 0: loss_avg = sum(loss_avg) / len(loss_avg) loss1_avg = sum(loss1_avg) / len(loss1_avg) loss2_avg = sum(loss2_avg) / len(loss2_avg) loss3_avg = sum(loss3_avg) / len(loss3_avg) t_end = time.time() time_interval = t_end - t_start logger.info( 'iter: {}, loss1: {:4f}, loss2: {:4f}, loss3:{:4f}, loss: {:4f}, lr: {:4f}, time: {:3f}' .format(count, loss1_avg, loss2_avg, loss3_avg, loss_avg, optim.lr, time_interval)) loss_avg = [] loss1_avg = [] loss2_avg = [] loss3_avg = [] t_start = t_end count += 1 if count == 25000: break ## dump model logger.info('saving trained model') torch.save(net.module.state_dict(), './res/model.pkl') torch.save(BNNeck.module.state_dict(), './res/BNNeck.pkl') logger.info('everything finished')
def __init__(self, cfg: CfgNode, num_classes): super().__init__() self.model = build_model(cfg, num_classes).cuda() # loss self.loss_type = cfg.LOSS.LOSS_TYPE self.loss_function_map = OrderedDict() # ID loss if 'softmax' in self.loss_type: self.xent = MyCrossEntropy( num_classes=num_classes, label_smooth=cfg.LOSS.IF_LABEL_SMOOTH, learning_weight=cfg.LOSS.IF_LEARNING_WEIGHT).cuda() def loss_function(data: Data): return cfg.LOSS.ID_LOSS_WEIGHT * self.xent( data.cls_score, data.cls_label) self.loss_function_map["softmax"] = loss_function if 'arcface' in self.loss_type: self.arcface = ArcfaceLoss(num_classes=num_classes, feat_dim=self.model.in_planes).cuda() def loss_function(data: Data): return self.arcface(data.feat_c, data.cls_label) self.loss_function_map["arcface"] = loss_function # metric loss if 'triplet' in self.loss_type: self.triplet = TripletLoss(cfg.LOSS.MARGIN, learning_weight=False).cuda() def loss_function(data: Data): return cfg.LOSS.METRIC_LOSS_WEIGHT * self.triplet( data.feat_t, data.cls_label) self.loss_function_map["triplet"] = loss_function # cluster loss if 'center' in self.loss_type: self.center = CenterLoss(num_classes=num_classes, feat_dim=self.model.in_planes, loss_weight=cfg.LOSS.CENTER_LOSS_WEIGHT, learning_weight=False).cuda() def loss_function(data: Data): return cfg.LOSS.CENTER_LOSS_WEIGHT * self.center( data.feat_t, data.cls_label) self.loss_function_map["center"] = loss_function self.optimizer = optim.Adam(self.model.parameters(), lr=cfg.OPTIMIZER.BASE_LR, weight_decay=cfg.OPTIMIZER.WEIGHT_DECAY) self.xent_optimizer = optim.SGD( self.xent.parameters(), lr=cfg.OPTIMIZER.LOSS_LR / 5, momentum=cfg.OPTIMIZER.MOMENTUM, weight_decay=cfg.OPTIMIZER.WEIGHT_DECAY) self.center_optimizer = optim.SGD( self.center.parameters(), lr=cfg.OPTIMIZER.LOSS_LR, momentum=cfg.OPTIMIZER.MOMENTUM, weight_decay=cfg.OPTIMIZER.WEIGHT_DECAY) model_list = [self.model, self.xent, self.triplet, self.center] opt_list = [self.optimizer, self.center_optimizer, self.xent_optimizer] model_list, opt_list = amp.initialize(model_list, opt_list, opt_level='O1', loss_scale=1.0) self.model, self.xent, self.triplet, self.center = model_list[ 0], model_list[1], model_list[2], model_list[3] self.optimizer, self.center_optimizer, self.xent_optimizer = opt_list[ 0], opt_list[1], opt_list[2] self.scheduler = WarmupMultiStepLR(self.optimizer, cfg.WARMUP.STEPS, cfg.WARMUP.GAMMA, cfg.WARMUP.FACTOR, cfg.WARMUP.MAX_EPOCHS, cfg.WARMUP.METHOD)
if use_triplet == True: train_data = DataLoader(train_datasets, sampler=RandomIdentitySampler_new(train, NUM_CLASSES ,4), batch_size=BATCH_SIZE, pin_memory=True, num_workers=8, drop_last=True) else: train_data = DataLoader(train_datasets, batch_size=BATCH_SIZE, pin_memory=True, shuffle=True) test_data = DataLoader(MyDataset(txt_path='/data/zhoumi/datasets/train_data/val.txt', transform=test_transform), batch_size=TEST_BATCH_SIZE, pin_memory=True) optimizer = optim.Adam(params=model.parameters(), lr=1e-4, weight_decay=5e-4) # optimizer = optim.SGD(model.parameters(), lr=1e-4, momentum=0.9, weight_decay=5e-4) #define loss function xent_criterion = CrossEntropyLabelSmooth(NUM_CLASSES) center_criterion = CenterLoss(NUM_CLASSES, feat_dim=1792) triplet_criterion = TripletLoss(margin=0.3) best_model = model best_acc = 0 print(len(test_data) * TEST_BATCH_SIZE, len(train_data)) model = model.cuda() for epoch in range(MAX_EPOC): lr = adjust_lr(epoch) for p in optimizer.param_groups: p['lr'] = lr for i, inputs in enumerate(train_data): model = model.train()
def main(): if args.dataset == "Office31": num_class = 31 elif args.dataset == "ImageCLEF": num_class = 12 else: raise ValueError("Unknown dataset {}".format(args.dataset)) if args.arch == "vgg16": model = DANNet(arch="vgg16", num_class=num_class) center_loss = CenterLoss(num_class, 2048).to('cuda') elif args.arch == "resnet50": model = DANNet(arch="resnet50", num_class=num_class) center_loss = CenterLoss(num_class, 2048).to('cuda') else: raise ValueError("Unknown model architecture {}".format(args.arch)) if len(args.gpu_ids) > 1: model = nn.DataParallel(model, device_ids=args.gpu_ids).cuda() else: model.cuda() # load source dataset default multi gpu if len(args.gpu_ids) > 1: transformer_train = model.module.train_augmentation() transformer_train_multi = model.module.train_multi_augmentation() transformer_test = model.module.test_augmentation() else: transformer_train = model.train_augmentation() transformer_train_multi = model.train_multi_augmentation() transformer_test = model.test_augmentation() source_train_loader = DataLoader(OfficeDataSet( data_path=args.train_path, transformer=transformer_train, k=1), batch_size=args.batch_size, num_workers=args.workers, shuffle=True, drop_last=True, pin_memory=False) # load target dataset for training, its label is not used for training default multi gpu target_train_loader = DataLoader(OfficeDataSet( data_path=args.test_path, transformer=transformer_train, k=args.k), batch_size=args.batch_size, num_workers=args.workers, shuffle=True, drop_last=True, pin_memory=False) # load target dataset for testing, its label is used for testing until the whole training process ends. target_test_loader = DataLoader(OfficeDataSet(data_path=args.test_path, transformer=transformer_test, k=1), batch_size=args.batch_size, num_workers=args.workers, shuffle=False, pin_memory=False) # set seed set_seed() # set writer to load some training information writer = SummaryWriter(args.tensorboard_file) logging.basicConfig(level=logging.INFO, filename=args.log_file, filemode='w') # set loss function criterion = nn.CrossEntropyLoss() # set optimizer if args.optimizer == "SGD": optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) optimizer_centloss = torch.optim.SGD(center_loss.parameters(), lr=0.5) elif args.optimizer == "Adam": optimizer = optim.Adam(model.parameters(), lr=args.lr) optimizer_centloss = torch.optim.SGD(center_loss.parameters(), lr=0.5) else: raise ValueError( "the author is too lazy to add the optimizer {}".format(args.opt)) # set lr scheduler if args.lr_scheduler == 'cosine_decay': lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, len(source_train_loader), eta_min=0) elif args.lr_scheduler == "step_decay": # warning! the milestones should be modified by your own demand lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[5, 10, 15], gamma=0.1) elif args.lr_scheduler == "custom": raise ValueError( "custom learning rate scheduler is under the todo list") else: raise ValueError("warning, wrong learning rate scheduler {}".format( args.lr_schedule)) # training process for epoch in trange(args.start_epoch, args.num_epoch): writer.add_scalar("train_epoch", epoch) logging.info("train_epoch: {}".format(epoch)) gamma = 2 / (1 + math.exp( -10 * (epoch) / args.num_epoch)) - 1 # set gamma for each epoch if args.method == "train": losses, mmdlosses, softmaxloss, train_acc = train( source_train_loader, target_train_loader, model, criterion, optimizer, epoch, gamma=gamma, k=args.k) elif args.method == "train_mixup": losses, mmdlosses, softmaxloss, train_acc = train_mixup( source_train_loader, target_train_loader, model, criterion, optimizer, epoch, gamma=gamma) elif args.method == "train_mixmatch": losses, softmaxloss, l2loss, mmdlosses, train_acc = train_mixmatch( source_train_loader, target_train_loader, model, criterion, optimizer, epoch, gamma=gamma, k=args.k) writer.add_scalar("l2_loss", l2loss) logging.info("l2_loss: {}".format(l2loss)) elif args.method == "consistency_regu": losses, softmaxloss, l2loss, mmdlosses, train_acc = train_consistency_regu( source_train_loader, target_train_loader, model, criterion, optimizer, epoch, gamma=gamma, k=args.k) elif args.method == "train_center": losses, mmdlosses, softmaxloss, train_acc = train_center( source_train_loader, target_train_loader, model, criterion, center_loss, optimizer, optimizer_centloss, epoch, gamma=gamma, k=args.k) else: raise ValueError("other tricks is under the todo list") # writer tensorboard value writer.add_scalar("train_loss", losses) writer.add_scalar("train_MMDloss", mmdlosses) writer.add_scalar("train_CEloss", softmaxloss) writer.add_scalar("train_acc", train_acc) # write result to log file logging.info("train_loss: {}".format(losses)) logging.info("train_MMDloss: {}".format(mmdlosses)) logging.info("train_CEloss: {}".format(softmaxloss)) logging.info("train_acc: {}".format(train_acc)) is_best = True save_checkpoint( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), }, is_best) if (epoch + 1) % args.eval_freq == 0 or epoch == args.num_epoch - 1: acc, val_loss = validate(target_test_loader, model, criterion) print("validation Epoch: {}, loss: {}, acc: {}".format( epoch + 1, val_loss, acc)) writer.add_scalar("val_epoch", epoch) writer.add_scalar("val_acc", acc) writer.add_scalar("val_loss", val_loss) logging.info("val_epoch: {}".format(epoch)) logging.info("val_acc: {}".format(acc)) logging.info("val_loss: {}".format(val_loss)) if args.lr_scheduler: lr_scheduler.step() # lr scheduler for each epoch print("Epoch: {} lr: {}".format(epoch + 1, lr_scheduler.get_lr()))