def get_trainloader_resnetAttr(dataset_reid,transform_train,pin_memory):
if args.dataset=='market1501' or args.dataset=='dukemtmcreid' or args.dataset=='pa100K':
trainloader_reid = DataLoader(
ImageDataset(dataset_reid.train,transform=transform_train,arch=args.arch),
sampler=RandomIdentitySampler(dataset_reid.train, args.train_batch, args.num_instances),
batch_size=args.train_batch, num_workers=args.workers,
pin_memory=pin_memory, drop_last=True,
)
return trainloader_reid
def main():
global args, best_rank1
torch.manual_seed(args.seed)
if not args.use_avai_gpus:
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,
)
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),
sampler=RandomIdentitySampler(dataset.train, args.train_batch,
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),
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),
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)))
if args.label_smooth:
criterion_xent = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
else:
criterion_xent = nn.CrossEntropyLoss()
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 and check_isfile(args.load_weights):
# load pretrained weights but ignore layers that don't match in size
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 and check_isfile(args.resume):
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
args.start_epoch = checkpoint['epoch'] + 1
best_rank1 = checkpoint['rank1']
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch,
best_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.visualize_ranks:
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_epoch = args.start_epoch
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 __init__(
self,
use_gpu,
source_names,
target_names,
root,
split_id=0,
height=256,
width=128,
train_batch_size=32,
test_batch_size=100,
workers=4,
train_sampler='',
data_augment='none',
num_instances=4, # number of instances per identity (for RandomIdentitySampler)
cuhk03_labeled=False, # use cuhk03's labeled or detected images
cuhk03_classic_split=False # use cuhk03's classic split or 767/700 split
):
super(ImageDataManager, self).__init__()
from torchreid.dataset_loader import ImageDataset
from torchreid.datasets import init_imgreid_dataset
from torchreid.transforms import build_transforms
from torch.utils.data import DataLoader
from torchreid.samplers import RandomIdentitySampler
self.use_gpu = use_gpu
self.source_names = source_names
self.target_names = target_names
self.root = root
self.split_id = split_id
self.height = height
self.width = width
self.train_batch_size = train_batch_size
self.test_batch_size = test_batch_size
self.workers = workers
self.train_sampler = train_sampler
self.num_instances = num_instances
self.cuhk03_labeled = cuhk03_labeled
self.cuhk03_classic_split = cuhk03_classic_split
self.pin_memory = True if self.use_gpu else False
# Build train and test transform functions
transform_train = build_transforms(self.height,
self.width,
is_train=True,
data_augment=data_augment)
transform_test = build_transforms(self.height,
self.width,
is_train=False,
data_augment=data_augment)
# transform_test_flip = build_transforms(self.height, self.width, is_train=False, data_augment=data_augment, flip=True)
print("=> Initializing TRAIN (source) datasets")
self.train = []
self._num_train_pids = 0
self._num_train_cams = 0
for name in self.source_names:
dataset = init_imgreid_dataset(
root=self.root,
name=name,
split_id=self.split_id,
cuhk03_labeled=self.cuhk03_labeled,
cuhk03_classic_split=self.cuhk03_classic_split)
for img_path, pid, camid in dataset.train:
pid += self._num_train_pids
camid += self._num_train_cams
self.train.append((img_path, pid, camid))
self._num_train_pids += dataset.num_train_pids
self._num_train_cams += dataset.num_train_cams
if self.train_sampler == 'RandomIdentitySampler':
print('!!! Using RandomIdentitySampler !!!')
self.trainloader = DataLoader(
ImageDataset(self.train, transform=transform_train),
sampler=RandomIdentitySampler(self.train,
self.train_batch_size,
self.num_instances),
batch_size=self.train_batch_size,
shuffle=False,
num_workers=self.workers,
pin_memory=self.pin_memory,
drop_last=True)
else:
self.trainloader = DataLoader(ImageDataset(
self.train, transform=transform_train),
batch_size=self.train_batch_size,
shuffle=True,
num_workers=self.workers,
pin_memory=self.pin_memory,
drop_last=True)
print("=> Initializing TEST (target) datasets")
self.testloader_dict = {
name: {
'query': None,
'gallery': None
}
for name in self.target_names
}
self.testdataset_dict = {
name: {
'query': None,
'gallery': None
}
for name in self.target_names
}
for name in self.target_names:
dataset = init_imgreid_dataset(
root=self.root,
name=name,
split_id=self.split_id,
cuhk03_labeled=self.cuhk03_labeled,
cuhk03_classic_split=self.cuhk03_classic_split)
self.testloader_dict[name]['new_vid_old_cid_query'] = DataLoader(
ImageDataset(dataset.new_vid_old_cid_query,
transform=transform_test),
batch_size=self.test_batch_size,
shuffle=False,
num_workers=self.workers,
pin_memory=self.pin_memory,
drop_last=False)
self.testloader_dict[name]['new_vid_old_cid_val'] = DataLoader(
ImageDataset(dataset.new_vid_old_cid_val,
transform=transform_test),
batch_size=self.test_batch_size,
shuffle=False,
num_workers=self.workers,
pin_memory=self.pin_memory,
drop_last=False)
self.testloader_dict[name]['new_vid_new_cid_query'] = DataLoader(
ImageDataset(dataset.new_vid_new_cid_query,
transform=transform_test),
batch_size=self.test_batch_size,
shuffle=False,
num_workers=self.workers,
pin_memory=self.pin_memory,
drop_last=False)
self.testloader_dict[name]['new_vid_new_cid_val'] = DataLoader(
ImageDataset(dataset.new_vid_new_cid_val,
transform=transform_test),
batch_size=self.test_batch_size,
shuffle=False,
num_workers=self.workers,
pin_memory=self.pin_memory,
drop_last=False)
self.testloader_dict[name]['train_gallery'] = DataLoader(
ImageDataset(dataset.train_gallery, transform=transform_test),
batch_size=self.test_batch_size,
shuffle=False,
num_workers=self.workers,
pin_memory=self.pin_memory,
drop_last=False)
# self.testdataset_dict[name]['query'] = dataset.query
# self.testdataset_dict[name]['gallery'] = dataset.gallery
print("\n")
print(" **************** Summary ****************")
print(" train names : {}".format(self.source_names))
print(" # train datasets : {}".format(len(self.source_names)))
print(" # train ids : {}".format(self._num_train_pids))
print(" # train images : {}".format(len(self.train)))
print(" # train cameras : {}".format(self._num_train_cams))
print(" test names : {}".format(self.target_names))
print(" *****************************************")
print("\n")
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, args.train_batch, 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 = nn.CrossEntropyLoss()
criterion_htri = TripletLoss(margin=args.margin)
criterion_KA = KALoss(margin=args.margin, same_margin = args.same_margin, use_auto_samemargin = args.use_auto_samemargin)
cirterion_lifted = LiftedLoss(margin=args.margin)
cirterion_batri = BA_TripletLoss(margin=args.margin)
if args.use_auto_samemargin == True:
G_params = [{'params': model.parameters(), 'lr': args.lr }, {'params': criterion_KA.auto_samemargin, 'lr': args.lr}]
else :
G_params = [para for _, para in model.named_parameters()]
optimizer = init_optim(args.optim, G_params, args.lr, args.weight_decay)
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, args.pool, 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()
adjust_learning_rate(optimizer, epoch)
train(epoch, model, cirterion_batri, cirterion_lifted, criterion_xent, criterion_htri, criterion_KA, 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))