def init_model(self):
print('Initializing model: {}'.format(args.model))
self.model = models.init_model(name=self.args.model,
num_classes=self.dm.num_attributes,
pretrained=not self.args.no_pretrained,
use_gpu=self.use_gpu)
print('Model size: {:.3f} M'.format(count_num_param(self.model)))
# Load pretrained weights if specified in args.
self.loaded_args = self.args
load_file = osp.join(args.save_experiment, args.load_weights)
if args.load_weights:
# TODO: implement result dict
if check_isfile(load_file):
cp = load_pretrained_weights([self.model], load_file)
if "args" in cp:
self.loaded_args = cp["args"]
else:
print("WARNING: Could not load args. ")
else:
print("WARNING: Could not load pretrained weights")
# Load model onto GPU if GPU is used.
self.model = self.model.cuda() if self.use_gpu else self.model
# Select Loss function.
if args.loss_func == "deepmar":
pos_ratio = self.dm.dataset.get_positive_attribute_ratio()
self.criterion = DeepMARLoss(pos_ratio,
args.train_batch_size,
use_gpu=self.use_gpu,
sigma=args.loss_func_param)
elif args.loss_func == "scel":
self.criterion = SigmoidCrossEntropyLoss(
num_classes=self.dm.num_attributes, use_gpu=self.use_gpu)
elif args.loss_func == "sscel":
attribute_grouping = self.dm.dataset.attribute_grouping
self.criterion = SplitSoftmaxCrossEntropyLoss(attribute_grouping,
use_gpu=self.use_gpu)
else:
self.criterion = None
self.f1_calibration_thresholds = None
self.optimizer = init_optimizer(self.model, **optimizer_kwargs(args))
self.scheduler = init_lr_scheduler(self.optimizer,
**lr_scheduler_kwargs(args))
self.model = nn.DataParallel(
self.model) if self.use_gpu else self.model
# Set default for max_epoch if it was not passed as an argument in the console.
if self.args.max_epoch < 0:
self.args.max_epoch = 30
self.model_list = [self.model]
self.optimizer_list = [self.optimizer]
self.scheduler_list = [self.scheduler]
self.criterion_list = [self.criterion]
def main():
torch.manual_seed(1)
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_devices
use_gpu = torch.cuda.is_available()
if use_cpu: use_gpu = False
if use_gpu:
print("Currently using GPU {}".format(gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(1)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing model: {}".format(arch))
model = init_model(name=arch, num_classes=576, loss_type=loss_type)
print("Model size: {:.3f} M".format(count_num_param(model)))
if use_gpu:
model = nn.DataParallel(model).cuda()
optimizer = init_optim(optim, model.parameters(), lr, weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=stepsize,
gamma=gamma)
model.train()
cnt = 0
for epoch in range(start_epoch, max_epoch, 2):
for step in range(2):
x = torch.randn(1, 3, 200, 200)
y = torch.randint(low=0, high=576, size=(1, ), dtype=torch.int64)
if use_gpu:
x = x.cuda()
y = y.cuda()
scheduler.step()
cnt += 1
print(cnt, scheduler.get_lr())
output = model(x)
# loss = nn.CrossEntropyLoss()(output[0], y)
loss = torch.tensor(0.0, dtype=torch.float32).cuda()
# loss = torch.tensor(0.0, dtype=torch.float32)
loss.requires_grad = True
optimizer.zero_grad()
loss.backward()
print(loss)
print(loss._grad)
optimizer.step()
print('Done.')
def main():
#GENERAL
root = "/home/kuru/Desktop/veri-gms-master/"
train_dir = '/home/kuru/Desktop/veri-gms-master/VeRispan/image_train/'
source = {'veri'}
target = {'veri'}
workers = 2
height = 320
width = 320
train_sampler = 'RandomSampler'
#AUGMENTATION
random_erase = True
jitter = True
aug = True
#OPTIMIZATION
opt = 'adam'
lr = 0.0003
weight_decay = 5e-4
momentum = 0.9
sgd_damp = 0.0
nesterov = True
warmup_factor = 0.01
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 = 6
lambda_tri = 1
#MODEL
arch = 'resnet101_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 = None
start_eval = 0
eval_freq = -1
#MISC
use_gpu = True
use_amp = True
print_freq = 50
seed = 1
resume = ''
save_dir = '/home/kuru/Desktop/veri-gms-master/logapex/'
gpu_id = 0
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')
train = []
num_train_pids = 0
num_train_cams = 0
for img_path, pid, camid in dataset.train:
path = img_path[52:77]
#print(path)
folder = path[1:4]
pid += num_train_pids
camid += num_train_cams
train.append((path, folder, pid, camid))
num_train_pids += dataset.num_train_pids
num_train_cams += dataset.num_train_cams
pid = 0
pidx = {}
for img_path, pid, camid in dataset.train:
path = img_path[52:77]
folder = path[1:4]
pidx[folder] = pid
pid += 1
path = '/home/kuru/Desktop/veri-gms-master/gms/'
pkl = {}
entries = os.listdir(path)
for name in entries:
f = open((path + name), 'rb')
if name == 'featureMatrix.pkl':
s = name[0:13]
else:
s = name[0:3]
pkl[s] = pickle.load(f)
f.close
transform_t = train_transforms(**transform_kwargs)
data_tfr = vd(
pkl_file='index.pkl',
dataset=train,
root_dir='/home/kuru/Desktop/veri-gms-master/VeRi/image_train/',
transform=transform_t)
trainloader = DataLoader(data_tfr,
sampler=None,
batch_size=train_batch_size,
shuffle=True,
num_workers=workers,
pin_memory=False,
drop_last=True)
#data_tfr = vd(pkl_file='index.pkl', dataset = train, root_dir=train_dir,transform=transforms.Compose([Rescale(64),RandomCrop(32),ToTensor()]))
#dataloader = DataLoader(data_tfr, batch_size=batch_size, shuffle=False, num_workers=0)
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'},
last_stride=1,
pretrained=not no_pretrained,
use_gpu=use_gpu)
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)
model = (model).cuda() if use_gpu else model
#model = nn.DataParallel(model).cuda() if use_gpu else model
optimizer = init_optimizer(model, **optimizer_kwargs)
#optimizer = init_optimizer(model)
model, optimizer = amp.initialize(model,
optimizer,
opt_level="O2",
keep_batchnorm_fp32=True,
loss_scale="dynamic")
model = nn.DataParallel(model).cuda() if use_gpu else model
scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs)
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)
print('=> Start training')
data_index = search(pkl)
for epoch in range(start, max_epoch):
losses = AverageMeter()
#xent_losses = AverageMeter()
htri_losses = AverageMeter()
accs = AverageMeter()
batch_time = 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) 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)
for i in range(train_batch_size):
labelx = str(label[i])
indexx = int(index[i])
cidx = int(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.argmax(ma.masked_where(a == 0, a))
pos_dic = data_tfr[data_index[cidx][1] + minpos]
neg_label = int(labelx)
while True:
neg_label = random.choice(range(1, 770))
if neg_label is not int(labelx) and os.path.isdir(
os.path.join(
'/home/kuru/Desktop/adiusb/veri-split/train',
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])))
neg_dic = data_tfr[data_index[neg_cid][1] + neg_index]
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]
#print("anc",labelx,'posdic', pos_dic[1],pos_dic[2],'negdic', neg_dic[1],neg_dic[2])
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)
#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
optimizer.zero_grad()
if use_amp:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
#loss.backward()
optimizer.step()
batch_time.update(time.time() - end)
losses.update(loss.item(), trainY.size(0))
htri_losses.update(htri_loss.item(), trainY.size(0))
accs.update(
accuracy(outputs[0:train_batch_size],
trainY[0:train_batch_size])[0])
if (batch_idx) % print_freq == 0:
print('Train ', end=" ")
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Acc {acc.val:.2f} ({acc.avg:.2f})\t'.format(
epoch + 1,
batch_idx + 1,
len(trainloader),
batch_time=batch_time,
loss=htri_losses,
acc=accs))
end = time.time()
scheduler.step()
print('=> Test')
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)
if vis_rank:
visualize_ranked_results(distmat,
dm.return_testdataset_by_name(name),
save_dir=osp.join(
save_dir, 'ranked_results', name),
topk=20)
def main():
torch.manual_seed(args.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu:
use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
train_dataset = data_manager.init_dataset(root=args.root,
name='json',
phase='train')
valid_dataset = data_manager.init_dataset(root=args.root,
name='json',
phase='valid')
test_dataset = data_manager.init_dataset(root=args.root,
name='json',
phase='test')
test_mask = test_dataset.mask
pin_memory = True if use_gpu else False
trainloader = DataLoader(JsonDataset(train_dataset),
num_workers=4,
batch_size=args.train_batch,
pin_memory=pin_memory,
drop_last=True)
validloader = DataLoader(JsonDataset(valid_dataset),
num_workers=4,
batch_size=args.test_batch,
pin_memory=pin_memory,
drop_last=True)
testloader = DataLoader(JsonDataset(test_dataset),
num_workers=4,
batch_size=args.test_batch,
pin_memory=pin_memory,
drop_last=True)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch, use_gpu=use_gpu)
print("Model size: {:.3f} M".format(count_num_param(model)))
model.init_weights()
criterion_label = LabelLoss()
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
start_epoch = args.start_epoch
if args.load_weights:
# load pretrained weights but ignore layers that don't match in size
print("Loading pretrained weights from '{}'".format(args.load_weights))
checkpoint = torch.load(args.load_weights)
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {
k: v
for k, v in pretrain_dict.items()
if k in model_dict and model_dict[k].size() == v.size()
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
if args.resume:
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}".format(start_epoch))
# if use_gpu:
# str_ids = args.gpu_devices.split(',')
# gpu_ids = []
# for str_id in str_ids:
# id = int(str_id)
# if id >= 0:
# gpu_ids.append(id)
# model = nn.DataParallel(model, gpu_ids)
device = torch.device(
'cuda' if use_gpu and torch.cuda.is_available() else 'cpu')
model.to(device)
if args.evaluate:
print("Evaluate only")
start_evaluate_time = time.time()
test_thetas = evaluate(model, testloader, use_gpu, args.test_batch,
test_mask)
# test_thetas = evaluate(model, validloader, use_gpu, args.test_batch, test_mask)
evaluate_time = time.time() - start_evaluate_time
print('Evaluate: {} secs'.format(evaluate_time))
with open("auto_sample.csv", "r") as csvfiler:
with open("test_thetas.csv", "w") as csvfilew:
reader = csv.reader(csvfiler)
for item in reader:
if reader.line_num == 1:
writer = csv.writer(csvfilew)
writer.writerow(['test_id', 'result'])
continue
writer = csv.writer(csvfilew)
writer.writerow(
[item[0],
str(test_thetas[reader.line_num - 2])])
# writer.writerows(map(lambda x: [x], test_thetas))
return
start_time = time.time()
train_time = 0
best_label_loss = np.inf
best_epoch = 0
#
# print("==> Test")
# label_loss = test(model, validloader, criterion_label, use_gpu, args.test_batch)
# print("test label loss RMES() is {}".format(label_loss))
#
# print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_label, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
# save model every epoch
if (epoch + 1) % args.save_step == 0:
print("==> Now save epoch {} \'s model".format(epoch + 1))
# if use_gpu:
# state_dict = model.state_dict() # module.
# else:
state_dict = model.state_dict()
save_checkpoint({
'state_dict': state_dict,
'epoch': epoch
}, False, osp.join(args.save_dir, 'checkpoint_latest.pth'))
# test model every eval_step
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or \
(epoch + 1) == args.max_epoch:
print("==> Test")
label_loss = test(model, validloader, criterion_label, use_gpu,
args.test_batch)
is_best = label_loss < best_label_loss
if is_best:
best_label_loss = label_loss
best_epoch = epoch + 1
# if use_gpu:
# state_dict = model.state_dict()
# else:
state_dict = model.state_dict()
save_checkpoint(
{
'state_dict': state_dict,
'label_loss': label_loss,
'epoch': epoch,
}, is_best,
osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) +
'.pth')) # .pth.tar
print("==> Best Label Loss {:.3}, achieved at epoch {}".format(
best_label_loss, 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))
def main():
torch.manual_seed(args.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_imgreid_dataset(
root=args.root, name=args.dataset, split_id=args.split_id,
cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split,
use_lmdb=args.use_lmdb,
)
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True if use_gpu else False
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train,
use_lmdb=args.use_lmdb, lmdb_path=dataset.train_lmdb_path),
batch_size=args.train_batch, shuffle=True, num_workers=args.workers,
pin_memory=pin_memory, drop_last=True,
)
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transform_test,
use_lmdb=args.use_lmdb, lmdb_path=dataset.query_lmdb_path),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test,
use_lmdb=args.use_lmdb, lmdb_path=dataset.gallery_lmdb_path),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent'}, use_gpu=use_gpu)
print("Model size: {:.3f} M".format(count_num_param(model)))
criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu)
optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma)
start_epoch = args.start_epoch
if args.fixbase_epoch > 0:
if hasattr(model, 'classifier') and isinstance(model.classifier, nn.Module):
optimizer_tmp = init_optim(args.optim, model.classifier.parameters(), args.fixbase_lr, args.weight_decay)
else:
print("Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0")
args.fixbase_epoch = 0
if args.load_weights:
# load pretrained weights but ignore layers that don't match in size
print("Loading pretrained weights from '{}'".format(args.load_weights))
checkpoint = torch.load(args.load_weights)
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
if args.resume:
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
rank1 = checkpoint['rank1']
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}\n- rank1: {}".format(start_epoch, rank1))
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
test(model, queryloader, galleryloader, use_gpu)
return
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print("==> Start training")
if args.fixbase_epoch > 0:
print("Train classifier for {} epochs while keeping base network frozen".format(args.fixbase_epoch))
for epoch in range(args.fixbase_epoch):
start_train_time = time.time()
train(epoch, model, criterion, optimizer_tmp, trainloader, use_gpu, freeze_bn=True)
train_time += round(time.time() - start_train_time)
del optimizer_tmp
print("Now open all layers for training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch+1) > args.start_eval and args.eval_step > 0 and (epoch+1) % args.eval_step == 0 or (epoch+1) == args.max_epoch:
print("==> Test")
rank1 = test(model, queryloader, galleryloader, 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(args.save_dir, 'checkpoint_ep' + str(epoch+1) + '.pth.tar'))
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))
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)
def main():
torch.manual_seed(args.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_vidreid_dataset(root=args.root,
name=args.dataset)
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True if use_gpu else False
# decompose tracklets into images for image-based training
new_train = []
for img_paths, pid, camid in dataset.train:
for img_path in img_paths:
new_train.append((img_path, pid, camid))
trainloader = DataLoader(
ImageDataset(new_train, transform=transform_train),
sampler=RandomIdentitySampler(new_train,
num_instances=args.num_instances),
batch_size=args.train_batch,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
queryloader = DataLoader(
VideoDataset(dataset.query,
seq_len=args.seq_len,
sample='evenly',
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.gallery,
seq_len=args.seq_len,
sample='evenly',
transform=transform_test),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'})
print("Model size: {:.3f} M".format(count_num_param(model)))
criterion_xent = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
start_epoch = args.start_epoch
if args.resume:
print("Loading checkpoint from '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
test(model, queryloader, galleryloader, args.pool, use_gpu)
return
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_xent, criterion_htri, optimizer,
trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (
epoch + 1) % args.eval_step == 0 or (epoch +
1) == args.max_epoch:
print("==> Test")
rank1 = test(model, queryloader, galleryloader, args.pool, 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(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
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))
def main():
#GENERAL
torch.cuda.empty_cache()
root = "/home/kuru/Desktop/veri-gms-master_noise/"
train_dir = '/home/kuru/Desktop/veri-gms-master_noise/VeRispan/image_train/'
source = {'verispan'}
target = {'verispan'}
workers = 4
height = 280
width = 280
train_size = 32
train_sampler = 'RandomSampler'
#AUGMENTATION
random_erase = True
jitter = True
aug = True
#OPTIMIZATION
opt = 'adam'
lr = 0.0003
weight_decay = 5e-4
momentum = 0.9
sgd_damp = 0.0
nesterov = True
warmup_factor = 0.01
warmup_method = 'linear'
#HYPERPARAMETER
max_epoch = 80
start = 0
train_batch_size = 8
test_batch_size = 100
#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='resnet101_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 = None
start_eval = 0
eval_freq = -1
#MISC
use_gpu = True
print_freq = 10
seed = 1
resume = ''
save_dir = '/home/kuru/Desktop/veri-gms-master_noise/spanningtree_verinoise_101_stride2/'
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_noise/', 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+6:90+6]
#print(path)
folder = path[1:4]
#print(folder)
pid += num_train_pids
newidd=0
train.append((path, folder, pid, camid,subid,countid))
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+6:90+6]
folder = path[1:4]
pidx[folder] = pid
pid+= 1
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
path = '/home/kuru/Desktop/adhi/veri-final-draft-master_noise/gmsNoise776/'
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))
with open('cids.pkl', 'rb') as handle:
b = pickle.load(handle)
#print(b)
with open('index.pkl', 'rb') as handle:
c = pickle.load(handle)
transform_t = train_transforms(**transform_kwargs)
data_tfr = vdspan(pkl_file='index_veryspan_noise.pkl', dataset = train, root_dir='/home/kuru/Desktop/veri-gms-master_noise/VeRispan/image_train/', transform=transform_t)
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
trainloader = DataLoader(data_tfr, sampler=None,batch_size=train_batch_size, shuffle=True, num_workers=workers,pin_memory=True, drop_last=True)
#data_tfr = vd(pkl_file='index.pkl', dataset = train, root_dir=train_dir,transform=transforms.Compose([Rescale(64),RandomCrop(32),ToTensor()]))
#dataloader = DataLoader(data_tfr, batch_size=batch_size, shuffle=False, num_workers=0)
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)
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)
scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs)
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)
print('=> Start training')
data_index = search(pkl)
print(len(data_index))
for epoch in range(start, max_epoch):
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)
for i in range(train_batch_size):
#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(minpos)
#print(np.array(data_index).shape)
#print(data_index[cidx][1])
pos_dic = data_tfr_old[data_index[cidx][1]+minpos]
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/adiusb/veri-split/train', 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])))
neg_dic = data_tfr_old[data_index[neg_cid][1]+neg_index]
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]
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)
#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
optimizer.zero_grad()
loss.backward()
optimizer.step()
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()
scheduler.step()
print('=> Test')
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)
#if (epoch + 1) == max_epoch:
if (epoch + 1) % 2 == 0:
print('=> Test')
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)
# if vis_rank:
# visualize_ranked_results(
# distmat, dm.return_testdataset_by_name(name),
# save_dir=osp.join(save_dir, 'ranked_results', name),
# topk=20)
save_checkpoint({
'state_dict': model.state_dict(),
'rank1': rank1,
'epoch': epoch + 1,
'arch': arch,
'optimizer': optimizer.state_dict(),
}, save_dir)
def main():
torch.manual_seed(args.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_imgreid_dataset(
root=args.root,
name=args.dataset,
split_id=args.split_id,
cuhk03_labeled=args.cuhk03_labeled,
cuhk03_classic_split=args.cuhk03_classic_split,
use_lmdb=args.use_lmdb,
)
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True if use_gpu else False
trainloader = DataLoader(
ImageDataset(dataset.train,
transform=transform_train,
use_lmdb=args.use_lmdb,
lmdb_path=dataset.train_lmdb_path),
sampler=RandomIdentitySampler(dataset.train,
num_instances=args.num_instances),
batch_size=args.train_batch,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
queryloader = DataLoader(
ImageDataset(dataset.query,
transform=transform_test,
use_lmdb=args.use_lmdb,
lmdb_path=dataset.train_lmdb_path),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery,
transform=transform_test,
use_lmdb=args.use_lmdb,
lmdb_path=dataset.train_lmdb_path),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'})
print("Model size: {:.3f} M".format(count_num_param(model)))
criterion_xent = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
if args.load_weights:
# load pretrained weights but ignore layers that don't match in size
if check_isfile(args.load_weights):
checkpoint = torch.load(args.load_weights)
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {
k: v
for k, v in pretrain_dict.items()
if k in model_dict and model_dict[k].size() == v.size()
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
print("Loaded pretrained weights from '{}'".format(
args.load_weights))
if args.resume:
if check_isfile(args.resume):
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
args.start_epoch = checkpoint['epoch']
rank1 = checkpoint['rank1']
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}\n- rank1: {}".format(
args.start_epoch, rank1))
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
distmat = test(model,
queryloader,
galleryloader,
use_gpu,
return_distmat=True)
if args.vis_ranked_res:
visualize_ranked_results(
distmat,
dataset,
save_dir=osp.join(args.save_dir, 'ranked_results'),
topk=20,
)
return
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print("==> Start training")
for epoch in range(args.start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_xent, criterion_htri, optimizer,
trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (
epoch + 1) % args.eval_step == 0 or (epoch +
1) == args.max_epoch:
print("==> Test")
rank1 = test(model, queryloader, galleryloader, 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(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
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))
def main():
torch.manual_seed(args.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_imgreid_dataset(
root=args.root,
name=args.dataset,
split_id=args.split_id,
cuhk03_labeled=args.cuhk03_labeled,
cuhk03_classic_split=args.cuhk03_classic_split,
use_lmdb=args.use_lmdb,
)
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
# T.Resize((args.height, args.width)),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True if use_gpu else False
trainset = ImageDataset(dataset.train,
transform=transform_train,
use_lmdb=args.use_lmdb,
lmdb_path=dataset.train_lmdb_path)
trainloader = DataLoader(
trainset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
queryset = ImageDataset(dataset.query,
transform=transform_test,
use_lmdb=args.use_lmdb,
lmdb_path=dataset.query_lmdb_path)
queryloader = DataLoader(
queryset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
galleryset = ImageDataset(dataset.gallery,
transform=transform_test,
use_lmdb=args.use_lmdb,
lmdb_path=dataset.gallery_lmdb_path)
galleryloader = DataLoader(
galleryset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'xent'},
use_gpu=use_gpu)
print("Model size: {:.3f} M".format(count_num_param(model)))
# summary(model, (3, 160, 64))
criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids,
use_gpu=use_gpu)
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
start_epoch = args.start_epoch
if args.fixbase_epoch > 0:
if hasattr(model, 'classifier') and isinstance(model.classifier,
nn.Module):
optimizer_tmp = init_optim(args.optim,
model.classifier.parameters(),
args.fixbase_lr, args.weight_decay)
else:
print(
"Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0"
)
args.fixbase_epoch = 0
if args.load_weights:
# load pretrained weights but ignore layers that don't match in size
print("Loading pretrained weights from '{}'".format(args.load_weights))
if torch.cuda.is_available():
checkpoint = torch.load(args.load_weights)
else:
checkpoint = torch.load(args.load_weights, map_location='cpu')
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {
k: v
for k, v in pretrain_dict.items()
if k in model_dict and model_dict[k].size() == v.size()
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
if args.resume:
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
rank1 = checkpoint['rank1']
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}\n- rank1: {}".format(start_epoch, rank1))
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
test(model, queryloader, galleryloader, use_gpu)
return
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print("==> Start training")
if args.fixbase_epoch > 0:
print(
"Train classifier for {} epochs while keeping base network frozen".
format(args.fixbase_epoch))
for epoch in range(args.fixbase_epoch):
start_train_time = time.time()
train(epoch,
model,
criterion,
optimizer_tmp,
trainloader,
use_gpu,
freeze_bn=True)
train_time += round(time.time() - start_train_time)
del optimizer_tmp
print("Now open all layers for training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (
epoch + 1) % args.eval_step == 0 or (epoch +
1) == args.max_epoch:
print("==> Test")
rank1 = test(model, queryloader, galleryloader, 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(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
'''
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint({
'state_dict': state_dict,
'epoch': epoch,
}, True, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch+1) + '.pth.tar'))
'''
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))
def main():
torch.manual_seed(1)
os.environ['CUDA_VISIBLE_DEVICES'] = config.gpu_devices
use_gpu = torch.cuda.is_available()
sys.stdout = Logger(config.save_dir, config.checkpoint_suffix,
config.evaluate)
print("\n==========\nArgs:")
config.print_parameter()
print("==========\n")
if use_gpu:
print("Currently using GPU {}".format(config.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(1)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(config.dataset))
dataset = data_manager.init_imgreid_dataset(name=config.dataset,
root=config.data_root)
transform_train = T.Compose([
T.Random2DTranslation(config.height, config.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=data_mean, std=data_std),
])
transform_test = T.Compose([
T.Resize((config.height, config.width)),
T.ToTensor(),
T.Normalize(mean=data_mean, std=data_std),
])
pin_memory = True if use_gpu else False
# train_batch_sampler = BalancedBatchSampler(dataset.train, n_classes=8, n_samples=8)
# train_batch_sampler = CCLBatchSampler(dataset.train, n_classes=n_classes, n_samples=n_samples)
# train_batch_sampler = CCLBatchSamplerV2(dataset.train, n_classes=n_classes, pos_samp_cnt=pos_samp_cnt,
# neg_samp_cnt=neg_samp_cnt, each_cls_max_cnt=each_cls_max_cnt)
train_batch_sampler = ClassSampler(dataset.train,
sample_cls_cnt=config.sample_cls_cnt,
each_cls_cnt=config.each_cls_cnt)
# trainloader = DataLoader(
# ImageDataset(dataset.train, transform=transform_train),
# batch_sampler=train_batch_sampler, batch_size=args.train_batch,
# shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True
# )
trainloader = DataLoader(ImageDatasetWCL(dataset,
data_type='train',
merge_h=256,
merge_w=256,
mean_std=[data_mean, data_std]),
batch_sampler=train_batch_sampler,
num_workers=config.workers,
pin_memory=pin_memory)
queryloader = DataLoader(
ImageDatasetWCL(dataset.query,
data_type='query',
merge_h=256,
merge_w=256,
mean_std=[data_mean, data_std]),
batch_size=config.test_batch,
shuffle=False,
num_workers=config.workers,
pin_memory=pin_memory,
drop_last=False,
)
galleryloader = DataLoader(
ImageDatasetWCL(dataset.gallery,
data_type='gallery',
merge_h=256,
merge_w=256,
mean_std=[data_mean, data_std]),
batch_size=config.test_batch,
shuffle=False,
num_workers=config.workers,
pin_memory=pin_memory,
drop_last=False,
)
if config.dataset == 'vehicleid':
train_query_loader = None
train_gallery_loader = None
else:
train_query_loader = DataLoader(
ImageDatasetWCL(dataset.train_query,
data_type='train_query',
merge_h=256,
merge_w=256,
mean_std=[data_mean, data_std]),
batch_size=config.test_batch,
shuffle=False,
num_workers=config.workers,
pin_memory=pin_memory,
drop_last=False,
)
train_gallery_loader = DataLoader(
ImageDatasetWCL(dataset.train_gallery,
data_type='train_gallery',
merge_h=256,
merge_w=256,
mean_std=[data_mean, data_std]),
batch_size=config.test_batch,
shuffle=False,
num_workers=config.workers,
pin_memory=pin_memory,
drop_last=False,
)
print("Initializing model: {}".format(config.arch))
model = init_model(name=config.arch,
num_classes=dataset.num_train_pids,
loss_type=config.loss_type)
print("Model size: {:.3f} M".format(count_num_param(model)))
if config.loss_type == 'xent':
criterion = [nn.CrossEntropyLoss(), nn.CrossEntropyLoss()]
elif config.loss_type == 'xent_triplet':
criterion = XentTripletLoss(
margin=config.margin,
triplet_selector=RandomNegativeTripletSelector(
margin=config.margin),
each_cls_cnt=config.each_cls_cnt,
n_class=config.sample_cls_cnt)
elif config.loss_type == 'xent_tripletv2':
criterion = XentTripletLossV2(
margin=config.margin,
triplet_selector=RandomNegativeTripletSelectorV2(
margin=config.margin),
each_cls_cnt=config.each_cls_cnt,
n_class=config.sample_cls_cnt)
# criterion = XentTripletLossV2(margin=0.04, triplet_selector=RandomNegativeTripletSelectorV2(margin=0.04),
# each_cls_cnt=config.each_cls_cnt, n_class=config.sample_cls_cnt)
# criterion = XentGroupTripletLossV2(margin=0.8, triplet_selector=AllTripletSelector(margin=0.8),
# each_cls_cnt=config.each_cls_cnt, n_class=config.sample_cls_cnt)
else:
raise KeyError("Unsupported loss: {}".format(config.loss_type))
optimizer = init_optim(config.optim, model.parameters(), config.lr,
config.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=config.stepsize,
gamma=config.gamma)
if config.resume is not None:
if check_isfile(config.resume):
checkpoint = torch.load(config.resume)
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {
k: v
for k, v in pretrain_dict.items()
if k in model_dict and model_dict[k].size() == v.size()
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
config.start_epoch = checkpoint['epoch']
rank1 = checkpoint['rank1']
if 'mAP' in checkpoint:
mAP = checkpoint['mAP']
else:
mAP = 0
print("Loaded checkpoint from '{}'".format(config.resume))
print("- start_epoch: {}\n- rank1: {}\n- mAP: {}".format(
config.start_epoch, rank1, mAP))
if use_gpu:
model = nn.DataParallel(model).cuda()
if config.evaluate:
print("Evaluate only")
test_model(model, queryloader, galleryloader, train_query_loader,
train_gallery_loader, use_gpu, config.test_batch,
config.loss_type, config.euclidean_distance_loss)
return
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_map = 0
best_epoch = 0
for epoch in range(config.start_epoch, config.max_epoch):
print("==> Start training")
start_train_time = time.time()
scheduler.step()
print('epoch:', epoch, 'lr:', scheduler.get_lr())
train(epoch, model, criterion, optimizer, trainloader,
config.loss_type, config.print_freq)
train_time += round(time.time() - start_train_time)
if epoch >= config.start_eval and config.eval_step > 0 and epoch % config.eval_step == 0 \
or epoch == config.max_epoch:
print("==> Test")
rank1, mAP = test_model(model, queryloader, galleryloader,
train_query_loader, train_gallery_loader,
use_gpu, config.test_batch,
config.loss_type,
config.euclidean_distance_loss)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_map = mAP
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,
'mAP': mAP,
'epoch': epoch + 1,
},
is_best,
use_gpu_suo=False,
fpath=osp.join(
config.save_dir, 'checkpoint_ep' + str(epoch + 1) +
config.checkpoint_suffix + '.pth.tar'))
print("==> Best Rank-1 {:.2%}, mAP {:.2%}, achieved at epoch {}".format(
best_rank1, best_map, 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))
def init_model(self):
print('Initializing main model: {}'.format(args.model))
self.model_main = models.init_model(
name=self.args.model,
num_classes=self.dm.num_attributes,
pretrained=not self.args.no_pretrained,
use_gpu=self.use_gpu)
print('Model size: {:.3f} M'.format(count_num_param(self.model_main)))
print('Initializing HP model: {}'.format(args.hp_model))
# Determine the size of the output vector for the HP-Net.
num_hp_net_outputs = 1 if self.args.hp_net_simple else self.dm.num_attributes
# Init the HP-Net
self.model_hp = models.init_model(
name="hp_net_" + self.args.hp_model,
num_classes=num_hp_net_outputs,
pretrained=not self.args.no_pretrained)
print('Model size: {:.3f} M'.format(count_num_param(self.model_hp)))
if self.args.rejector == "none":
self.rejector = rejectors.NoneRejector()
elif self.args.rejector == 'macc':
self.rejector = rejectors.MeanAccuracyRejector(
self.args.max_rejection_quantile)
elif self.args.rejector == "median":
self.rejector = rejectors.MedianRejector(
self.args.max_rejection_quantile)
elif self.args.rejector == "threshold":
self.rejector = rejectors.ThresholdRejector(
self.args.rejection_threshold,
self.args.max_rejection_quantile)
elif self.args.rejector == "quantile":
self.rejector = rejectors.QuantileRejector(
self.args.max_rejection_quantile)
elif self.args.rejector == 'f1':
self.rejector = rejectors.F1Rejector(
self.args.max_rejection_quantile)
else:
raise ValueError("Unsupported rejection strategy: '{}'".format(
self.args.rejector))
print("Using rejection strategy '{}'".format(self.args.rejector))
if self.args.hp_calib == 'none':
self.hp_calibrator = NoneCalibrator()
elif self.args.hp_calib == 'linear':
self.hp_calibrator = LinearCalibrator()
else:
raise ValueError("Unsupported calibrator: '{}'".format(
self.args.hp_calib))
print("Using calibrator for HP-Loss '{}'".format(self.args.hp_calib))
# Load pretrained weights if specified in args.
load_file = osp.join(args.save_experiment, args.load_weights)
self.loaded_args = self.args
if args.load_weights:
if check_isfile(load_file):
cp = load_pretrained_weights([self.model_main, self.model_hp],
load_file)
if "args" in cp:
self.loaded_args = cp["args"]
else:
print("WARNING: Could not load args. ")
if "result_dict" in cp and cp[
"result_dict"] is not None and self.args.evaluate:
self.result_dict = cp["result_dict"]
self.result_manager = ResultManager(self.result_dict)
print("Loaded result dict with keys: ")
print(sorted(list(self.result_dict.keys())))
if "rejection_thresholds" in self.result_dict:
self.rejector.load_thresholds(
self.result_dict["rejection_thresholds"])
if self.rejector.is_initialized():
print("Loaded rejection thresholds. ")
else:
print(
"Loaded uninitialized (None) rejection thresholds. "
)
else:
print("WARNING: Could not load rejection thresholds. ")
else:
print("WARNING: Could not load pretrained weights")
self.new_eval_split = self.args.eval_split != self.loaded_args.eval_split
# Load model onto GPU if GPU is used.
self.model_main = self.model_main.cuda(
) if self.use_gpu else self.model_main
self.model = self.model_main
self.model_hp = self.model_hp.cuda() if self.use_gpu else self.model_hp
self.pos_ratio = self.dm.dataset.get_positive_attribute_ratio()
# Select Loss function.
# Select Loss function.
if args.loss_func == "deepmar":
self.criterion = DeepMARLoss(self.pos_ratio,
args.train_batch_size,
use_gpu=self.use_gpu,
sigma=args.loss_func_param)
elif args.loss_func == "scel":
self.criterion = SigmoidCrossEntropyLoss(
num_classes=self.dm.num_attributes, use_gpu=self.use_gpu)
else:
self.criterion = None
self.criterion_main = self.criterion
self.criterion_hp = HardnessPredictorLoss(
self.args.use_deepmar_for_hp,
self.pos_ratio,
self.dm.num_attributes,
use_gpu=self.use_gpu,
sigma=self.args.hp_loss_param,
use_visibility=self.args.use_bbs_gt,
visibility_weight=self.args.hp_visibility_weight)
self.f1_calibration_thresholds = None
self.optimizer_main = init_optimizer(self.model_main,
**optimizer_kwargs(args))
self.scheduler_main = init_lr_scheduler(self.optimizer_main,
**lr_scheduler_kwargs(args))
self.optimizer = self.optimizer_main
self.scheduler = self.scheduler_main
op_args = optimizer_kwargs(args)
sc_args = lr_scheduler_kwargs(args)
op_args['lr'] *= self.args.hp_net_lr_multiplier
self.optimizer_hp = init_optimizer(self.model_hp, **op_args)
sc_args["stepsize"] = [
i + self.args.hp_epoch_offset for i in sc_args["stepsize"]
]
self.scheduler_hp = init_lr_scheduler(self.optimizer_hp, **sc_args)
self.model_main = nn.DataParallel(
self.model_main) if self.use_gpu else self.model_main
self.model = self.model_main
self.model_hp = nn.DataParallel(
self.model_hp) if self.use_gpu else self.model_hp
if not self.args.evaluate:
self.init_epochs()
self.model_list = [self.model_main, self.model_hp]
self.optimizer_list = [self.optimizer_main, self.optimizer_hp]
self.scheduler_list = [self.scheduler_main, self.scheduler_hp]
self.criterion_list = [self.criterion_main, self.criterion_hp]
def main():
torch.manual_seed(args.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_vidreid_dataset(root=args.root, name=args.dataset)
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True if use_gpu else False
# decompose tracklets into images for image-based training
new_train = []
for img_paths, pid, camid in dataset.train:
for img_path in img_paths:
new_train.append((img_path, pid, camid))
trainloader = DataLoader(
ImageDataset(new_train, transform=transform_train),
sampler=RandomIdentitySampler(new_train, num_instances=args.num_instances),
batch_size=args.train_batch, num_workers=args.workers,
pin_memory=pin_memory, drop_last=True,
)
queryloader = DataLoader(
VideoDataset(dataset.query, seq_len=args.seq_len, sample='evenly', 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.gallery, seq_len=args.seq_len, sample='evenly', transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'})
print("Model size: {:.3f} M".format(count_num_param(model)))
criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma)
start_epoch = args.start_epoch
if args.load_weights:
# load pretrained weights but ignore layers that don't match in size
print("Loading pretrained weights from '{}'".format(args.load_weights))
checkpoint = torch.load(args.load_weights)
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
if args.resume:
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
rank1 = checkpoint['rank1']
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}\n- rank1: {}".format(start_epoch, rank1))
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
test(model, queryloader, galleryloader, args.pool, use_gpu)
return
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch+1) > args.start_eval and args.eval_step > 0 and (epoch+1) % args.eval_step == 0 or (epoch+1) == args.max_epoch:
print("==> Test")
rank1 = test(model, queryloader, galleryloader, args.pool, 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(args.save_dir, 'checkpoint_ep' + str(epoch+1) + '.pth.tar'))
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))
