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