def main(check_model, mm=1):
dataset = data_manager.init_img_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]),
])
'''trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train),
batch_size=args.train_batch, shuffle=True, num_workers=args.workers,
pin_memory=pin_memory, drop_last=True,
)'''
query = ImageDataset(dataset.query, transform=transform_test)
if args.dataset == 'beijing':
query = ImageDataset_forBeijing(dataset.query,
transform=transform_test)
#gallery = ImageDatasetLazy(dataset.gallery, transform=transform_test)
gallery = ImageDataset(dataset.gallery, transform=transform_test)
if args.dataset == 'beijing':
gallery = ImageDataset_forBeijing(dataset.gallery,
transform=transform_test)
if args.evaluate:
#print("Evaluate only")
if mm == 1:
cost, recall, precision = test(query, gallery, check_model, mm)
return cost, recall, precision
else:
cost, recall, precision, delay = test(query, gallery, check_model,
mm)
return cost, recall, precision, delay
def init_data_loaders(args, use_gpu=True):
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_dataset(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
trainloader = DataLoader(
VideoDataset(dataset.train, seq_len=args.seq_len, sample='random',transform=transform_train),
sampler=RandomIdentitySampler(dataset.train, num_instances=args.num_instances),
batch_size=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='random', 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='random', transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
return dataset, trainloader, queryloader, galleryloader
def train_each_teacher(num_epoch, train_data, train_label, test_data,
test_label, save_path):
torch.manual_seed(config.seed)
print('len of train_data in network', len(train_data))
os.environ['CUDA_VISIBLE_DEVICES'] = config.gpu_devices
print('it is training now')
use_gpu = torch.cuda.is_available()
if config.use_cpu: use_gpu = False
print('whether evaluate', config.evaluate)
if use_gpu:
print("Currently using GPU {}".format(config.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(config.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
transform_train = T.Compose([
T.Random2DTranslation(config.height, config.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((config.height, config.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(train_data, label=train_label, transform=transform_train),
batch_size=config.train_batch,
shuffle=True,
num_workers=config.workers,
pin_memory=pin_memory,
drop_last=False,
)
testloader = DataLoader(
ImageDataset(test_data, label=test_label, transform=transform_test),
batch_size=config.test_batch,
shuffle=False,
num_workers=config.workers,
pin_memory=pin_memory,
drop_last=False,
)
print("Initializing model: {}".format('resnet50m'))
model = models.init_model(name=config.arch,
num_classes=config.nb_labels,
loss={'xent'},
use_gpu=use_gpu)
if use_gpu:
model = nn.DataParallel(model).cuda()
criterion = nn.MultiLabelSoftMarginLoss()
optimizer = init_optim(config.optim, model.parameters(), config.lr,
config.weight_decay)
if config.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=config.stepsize,
gamma=config.gamma)
print("==> Start training")
start_time = time.time()
for epoch in range(num_epoch):
train(epoch, model, criterion, optimizer, trainloader, use_gpu)
if config.stepsize > 0: scheduler.step()
rank1 = test(model, testloader, use_gpu)
rank1 = test(model, testloader, use_gpu)
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
print('save model', save_path)
torch.save(state_dict, save_path)
#print("==> Hamming Score {:.3%}".format(rank1))
elapsed = round(time.time() - start_time)
print("Finished. Training time (h:m:s): {}.".format(elapsed))
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
print('whether evaluate', args.evaluate)
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_img_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
trainloader = DataLoader(
ImageDataset(dataset.train_data,
dataset.train_label,
transform=transform_train),
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
testloader = DataLoader(
ImageDataset(dataset.test_data,
dataset.test_label,
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=40,
loss={'xent'},
use_gpu=use_gpu)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion = nn.MultiLabelSoftMarginLoss()
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=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
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:
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, testloader, 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, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
if args.stepsize > 0: 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, testloader, 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,
'hamming_score': rank1,
'epoch': epoch,
}, is_best,
osp.join(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
print("==> Best Hamming Score {:.3%}, 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():
args.save_dir = args.save_dir + '/' + args.arch
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
# add data to save_dir
args.save_dir = args.save_dir + '_' + args.dataset + '_combined_multisteplr11'
if args.pretrained_model is not None:
args.save_dir = os.path.dirname(args.pretrained_model)
if not osp.exists(args.save_dir):
os.makedirs(args.save_dir)
log_name = 'test.log' if args.evaluate else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(args.save_dir, log_name))
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_dataset(name=args.dataset)
print("Train Transforms: \n\
Random2DTranslation, \n\
RandomHorizontalFlip, \n\
ToTensor, \n\
normalize\
")
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
# T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
# T.RandomErasing(p=0.5, scale=(0.02, 0.4), ratio=(0.3, 3.3), value=[0.485, 0.456, 0.406])
])
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(
VideoDataset(dataset.train, seq_len=args.seq_len,
sample=args.data_selection, transform=transform_train),
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(
VideoDataset(dataset.query, seq_len=args.seq_len,
sample='dense', 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='dense', 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, seq_len=args.seq_len)
# pretrained model loading
if args.pretrained_model is not None:
if not os.path.exists(args.pretrained_model):
raise IOError("Can't find pretrained model: {}".format(
args.pretrained_model))
print("Loading checkpoint from '{}'".format(args.pretrained_model))
pretrained_state = torch.load(args.pretrained_model)['state_dict']
print(len(pretrained_state), ' keys in pretrained model')
current_model_state = model.state_dict()
pretrained_state = {key: val
for key, val in pretrained_state.items()
if key in current_model_state and val.size() == current_model_state[key].size()}
print(len(pretrained_state),
' keys in pretrained model are available in current model')
current_model_state.update(pretrained_state)
model.load_state_dict(current_model_state)
print("Model size: {:.5f}M".format(sum(p.numel()
for p in model.parameters())/1000000.0))
if use_gpu:
model = nn.DataParallel(model).cuda()
criterion_xent = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = torch.optim.Adam(
model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(
optimizer, step_size=args.stepsize, gamma=args.gamma)
start_epoch = args.start_epoch
if args.evaluate:
print("Evaluate only")
test(model, queryloader, galleryloader, use_gpu)
return
start_time = time.time()
best_rank1 = -np.inf
is_first_time = True
for epoch in range(start_epoch, args.max_epoch):
eta_seconds = (time.time() - start_time) * (args.max_epoch - epoch) / max(epoch, 1)
eta_str = str(datetime.timedelta(seconds=int(eta_seconds)))
print("==> Epoch {}/{} \teta {}".format(epoch+1, args.max_epoch, eta_str))
train(model, criterion_xent, criterion_htri,
optimizer, trainloader, use_gpu)
if args.stepsize > 0:
scheduler.step()
rank1 = 'NA'
mAP = 'NA'
is_best = False
if args.eval_step > 0 and (epoch+1) % args.eval_step == 0 or (epoch+1) == args.max_epoch:
print("==> Test")
rank1, mAP = test(model, queryloader, galleryloader, use_gpu)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
# save the model as required
if (epoch+1) % args.save_step == 0:
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, args.save_prefix, 'model' + '.pth.tar-' + str(epoch+1)))
is_first_time = False
if not is_first_time:
utils.disable_all_print_once()
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
def attr_main():
runId = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
args.save_dir = os.path.join(args.save_dir, runId)
if not os.path.exists(args.save_dir):
os.mkdir(args.save_dir)
print(args.save_dir)
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('./log_train_' + runId + '.txt')
else:
sys.stdout = Logger('./log_test_' + runId + '.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_dataset(name=args.dataset,
min_seq_len=args.seq_len,
attr=True)
args.attr_lens = dataset.attr_lens
args.columns = dataset.columns
print("Initializing model: {}".format(args.arch))
# if args.arch == "resnet50ta_attr" or args.arch == "resnet50ta_attr_newarch":
if args.arch == 'attr_resnet503d':
model = models.init_model(name=args.arch,
attr_lens=args.attr_lens,
model_type=args.model_type,
num_classes=dataset.num_train_pids,
sample_width=args.width,
sample_height=args.height,
sample_duration=args.seq_len)
torch.backends.cudnn.benchmark = False
else:
model = models.init_model(name=args.arch,
attr_lens=args.attr_lens,
model_type=args.model_type)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
if args.dataset == "duke":
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])
])
elif args.dataset == "mars":
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(
VideoDataset(dataset.train,
seq_len=args.seq_len,
sample='random',
transform=transform_train,
attr=True,
attr_loss=args.attr_loss,
attr_lens=args.attr_lens),
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 = VideoDataset(dataset.query + dataset.gallery,
seq_len=args.seq_len,
sample='dense',
transform=transform_test,
attr=True,
attr_loss=args.attr_loss,
attr_lens=args.attr_lens)
start_epoch = args.start_epoch
if use_gpu:
model = nn.DataParallel(model).cuda()
start_time = time.time()
if args.arch == 'resnet503d':
torch.backends.cudnn.benchmark = False
# print("Run attribute pre-training")
if args.attr_loss == "cropy":
criterion = nn.CrossEntropyLoss()
elif args.attr_loss == "mse":
criterion = nn.MSELoss()
if args.evaluate:
print("Evaluate only")
model_root = "/data/chenzy/models/mars/2019-02-26_21-02-13"
model_paths = []
for m in os.listdir(model_root):
if m.endswith("pth"):
model_paths.append(m)
model_paths = sorted(model_paths,
key=lambda a: float(a.split("_")[1]),
reverse=True)
# model_paths = ['rank1_2.8755379380596713_checkpoint_ep500.pth']
for m in model_paths:
model_path = os.path.join(model_root, m)
print(model_path)
old_weights = torch.load(model_path)
new_weights = model.module.state_dict()
for k in new_weights:
if k in old_weights:
new_weights[k] = old_weights[k]
model.module.load_state_dict(new_weights)
avr_acc = attr_test(model, criterion, queryloader, use_gpu)
# break
# test(model, queryloader, galleryloader, args.pool, use_gpu)
return
if use_gpu:
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.module.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
# avr_acc = attr_test(model, criterion, queryloader, use_gpu)
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=args.stepsize,
gamma=args.gamma)
best_avr = 0
no_rise = 0
for epoch in range(start_epoch, args.max_epoch):
print("==> Epoch {}/{}".format(epoch + 1, args.max_epoch))
attr_train(model, criterion, optimizer, trainloader, use_gpu)
if args.stepsize > 0: scheduler.step()
if args.eval_step > 0 and ((epoch + 1) % (args.eval_step) == 0 or
(epoch + 1) == args.max_epoch):
avr_acc = attr_test(model, criterion, queryloader, use_gpu)
print("avr", avr_acc)
if avr_acc > best_avr:
no_rise = 0
print("==> Test")
best_avr = avr_acc
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
torch.save(
state_dict,
osp.join(
args.save_dir, "avr_" + str(avr_acc) +
'_checkpoint_ep' + str(epoch + 1) + '.pth'))
else:
no_rise += 1
print("no_rise:", no_rise)
if no_rise > 20:
break
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
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_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),
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),
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={'cent'})
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion_xent = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
criterion_cent = CenterLoss(num_classes=dataset.num_train_pids,
feat_dim=model.feat_dim,
use_gpu=use_gpu)
optimizer_model = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
optimizer_cent = torch.optim.SGD(criterion_cent.parameters(),
lr=args.lr_cent)
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer_model,
step_size=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, 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_cent, optimizer_model,
optimizer_cent, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
if args.stepsize > 0: scheduler.step()
if 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_dataset(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
trainloader = DataLoader(
VideoDataset(dataset.train,
seq_len=args.seq_len,
sample='random',
transform=transform_train),
sampler=RandomIdentitySampler(dataset.train,
num_instances=args.num_instances),
batch_size=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='dense',
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='dense',
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))
if args.arch == 'resnet503d':
model = resnet3d.resnet50(num_classes=dataset.num_train_pids,
sample_width=args.width,
sample_height=args.height,
sample_duration=args.seq_len)
if not os.path.exists(args.pretrained_model):
raise IOError("Can't find pretrained model: {}".format(
args.pretrained_model))
print("Loading checkpoint from '{}'".format(args.pretrained_model))
checkpoint = torch.load(args.pretrained_model)
state_dict = {}
for key in checkpoint['state_dict']:
if 'fc' in key: continue
state_dict[key.partition("module.")
[2]] = checkpoint['state_dict'][key]
model.load_state_dict(state_dict, strict=False)
else:
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'})
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion_xent = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=args.stepsize,
gamma=args.gamma)
start_epoch = args.start_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()
best_rank1 = -np.inf
if args.arch == 'resnet503d':
torch.backends.cudnn.benchmark = False
for epoch in range(start_epoch, args.max_epoch):
print("==> Epoch {}/{}".format(epoch + 1, args.max_epoch))
train(model, criterion_xent, criterion_htri, optimizer, trainloader,
use_gpu)
if args.stepsize > 0: scheduler.step()
if 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
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'))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
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_dataset(root=args.root,
name=args.dataset,
cls_sample=args.cls_sample)
# print(dataset.train)
# print(1)
# 解释器：创建一个transform处理图像数据的设置
# T.Random2DTranslation：随机裁剪
# T.RandomHorizontalFlip: 给定概率进行随机水平翻转
# T.ToTensor: 将PIL或numpy向量[0,255]=>tensor[0.0,1.0]
# T.Normalize：用均值和标准偏差标准化张量图像，mean[ , , ]三个参数代表三通道
transform_train = T.Compose([
# T.RandomCrop(224),
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(256),
# T.CenterCrop(224),
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
m = dataset.train
print(1)
# Dataloader 提供队列和线程
# ImageDataset:return data =>img, pid, camid
# RandomIdentitySampler:定义从数据集中抽取样本的策略
# num_workers: 子进程数
# print(dataset.train)
trainloader = DataLoader(
AGE_Gender_ImageDataset(dataset.train, transform=transform_train),
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
testloader = DataLoader(
AGE_Gender_ImageDataset(dataset.test, 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.train_num_class,
loss={'xent'})
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
# criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.train_num_class, use_gpu=use_gpu)
age_criterion_xent = nn.CrossEntropyLoss()
gender_criterion_xent = nn.CrossEntropyLoss()
# optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=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, testloader, use_gpu)
return
start_time = time.time()
train_time = 0
# best_rank1 = -np.inf
best_score = 0
best_MAE = 0
best_gender_acc = 0
best_epoch = 0
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, age_criterion_xent, gender_criterion_xent,
optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
if args.stepsize > 0: scheduler.step()
if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (
epoch + 1) == args.max_epoch:
print("==> Test")
MAE, Gender_acc = test(model, testloader, use_gpu)
Score = Gender_acc * 100 - MAE
is_best = Score > best_score
if is_best:
best_score = Score
best_MAE = MAE
best_gender_acc = Gender_acc
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': Score,
'epoch': epoch,
}, is_best,
osp.join(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
print(
"==> Best best_score(Gender_acc-MAE) {} |Gender_acc {}\t MAE {}|achieved at epoch {}"
.format(best_score, best_gender_acc, best_MAE, 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_dataset(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 = False
trainloader = DataLoader(
VideoDataset(dataset.train,
seq_len=args.seq_len,
sample='random',
transform=transform_train),
sampler=RandomIdentitySampler(dataset.train,
num_instances=args.num_instances),
batch_size=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='dense',
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='dense',
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: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion_xent = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=args.stepsize,
gamma=args.gamma)
start_epoch = args.start_epoch
start_time = time.time()
print(start_time)
for batch_idx, (imgs, pids, _) in enumerate(trainloader):
print(batch_idx)
print('x')
if use_gpu:
imgs, pids = imgs.cuda(), pids.cuda()
imgs, pids = Variable(imgs), Variable(pids)
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
def testseq(dataset_name, use_gpu):
dataset_root = './video2img/track1_sct_img_test_big/'
dataset = Graph_data_manager.AICityTrack2(root=dataset_root)
width = 224
height = 224
transform_train = T.Compose([
T.Random2DTranslation(height, 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((height, 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
seq_len = 4
num_instance = 4
train_batch = 32
test_batch = 1
queryloader = DataLoader(
VideoDataset(dataset.query,
seq_len=seq_len,
sample='dense',
transform=transform_test),
batch_size=test_batch,
shuffle=False,
num_workers=4,
pin_memory=pin_memory,
drop_last=False,
)
arch = "resnet50ta"
pretrained_model = "./log/track12_ta224_checkpoint_ep500.pth.tar"
start_epoch = 0
print("Initializing model: {}".format(arch))
dataset.num_train_pids = 517
if arch == 'resnet503d':
model = resnet3d.resnet50(num_classes=dataset.num_train_pids,
sample_width=width,
sample_height=height,
sample_duration=seq_len)
if not os.path.exists(pretrained_model):
raise IOError(
"Can't find pretrained model: {}".format(pretrained_model))
print("Loading checkpoint from '{}'".format(pretrained_model))
checkpoint = torch.load(pretrained_model)
state_dict = {}
for key in checkpoint['state_dict']:
if 'fc' in key: continue
state_dict[key.partition("module.")
[2]] = checkpoint['state_dict'][key]
model.load_state_dict(state_dict, strict=False)
else:
if not os.path.exists(pretrained_model):
model = models.init_model(name=arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'})
else:
model = models.init_model(name=arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'})
checkpoint = torch.load(pretrained_model)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch'] + 1
print("Loaded checkpoint from '{}'".format(pretrained_model))
print("- start_epoch: {}\n- rank1: {}".format(
start_epoch, checkpoint['rank1']))
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion_xent = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
criterion_htri = TripletLoss(margin=0.3)
lr = 0.0003
gamma = 0.1
stepsize = 200
weight_decay = 5e-04
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
weight_decay=weight_decay)
if stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=stepsize,
gamma=gamma)
start_epoch = start_epoch
if use_gpu:
model = nn.DataParallel(model).cuda()
test(model, queryloader, 'avg', use_gpu, dataset, -1, meta_data_tab=None)
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_dataset(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_train_p = T.Compose([
T.Random2DTranslation(256, 128),
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]),
])
transform_test_p = T.Compose([
T.Resize((256, 128)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_file = 'data/cuhk_train.pkl'
test_file = 'data/cuhk_test.pkl'
gallery_file = 'data/cuhk_gallery.pkl'
data_root = args.data_root
dataset_train = CUHKGroup(train_file, data_root, True, transform_train,
transform_train_p)
dataset_test = CUHKGroup(test_file, data_root, False, transform_test,
transform_test_p)
dataset_query = CUHKGroup(test_file, data_root, False, transform_test,
transform_test_p)
dataset_gallery = CUHKGroup(gallery_file, data_root, False, transform_test,
transform_test_p)
pin_memory = True if use_gpu else False
if args.xent_only:
trainloader = DataLoader(
dataset_train,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
else:
trainloader = DataLoader(
dataset_train,
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(
dataset_test,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
querygalleryloader = DataLoader(
dataset_query,
batch_size=args.gallery_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
galleryloader = DataLoader(
dataset_gallery,
batch_size=args.gallery_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
print("Initializing model: {}".format(args.arch))
if args.xent_only:
# model = models.init_model(name=args.arch, num_classes=dataset_train.num_train_gids, loss={'xent'})
model = models.init_model(name=args.arch,
num_classes=dataset_train.num_train_gids,
loss={'xent'})
else:
# model = models.init_model(name=args.arch, num_classes=dataset_train.num_train_gids, loss={'xent', 'htri'})
model = models.init_model(
name=args.arch,
num_classes=dataset_train.num_train_gids,
num_person_classes=dataset_train.num_train_pids,
loss={'xent', 'htri'})
#criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset_train.num_train_gids, use_gpu=use_gpu)
#criterion_xent_person = CrossEntropyLabelSmooth(num_classes=dataset_train.num_train_pids, use_gpu=use_gpu)
if os.path.exists(args.pretrained_model):
print("Loading checkpoint from '{}'".format(args.pretrained_model))
checkpoint = torch.load(args.pretrained_model)
model_dict = model.state_dict()
pretrain_dict = checkpoint['state_dict']
pretrain_dict = {
k: v
for k, v in pretrain_dict.items() if k in model_dict
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
criterion_xent = nn.CrossEntropyLoss(ignore_index=-1)
criterion_xent_person = nn.CrossEntropyLoss(ignore_index=-1)
criterion_htri = TripletLoss(margin=args.margin)
criterion_pair = ContrastiveLoss(margin=args.margin)
criterion_htri_filter = TripletLossFilter(margin=args.margin)
criterion_permutation = PermutationLoss()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
if args.stepsize > 0:
if args.warmup:
scheduler = WarmupMultiStepLR(optimizer, [200, 400, 600])
else:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=args.stepsize,
gamma=args.gamma)
start_epoch = args.start_epoch
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
test_gcn_person_batch(model, queryloader, querygalleryloader,
galleryloader, args.pool, use_gpu)
#test_gcn_batch(model, queryloader, querygalleryloader, galleryloader, args.pool, use_gpu)
#test_gcn(model, queryloader, galleryloader, args.pool, use_gpu)
#test(model, queryloader, galleryloader, args.pool, use_gpu)
return
start_time = time.time()
best_rank1 = -np.inf
for epoch in range(start_epoch, args.max_epoch):
#print("==> Epoch {}/{} lr:{}".format(epoch + 1, args.max_epoch, scheduler.get_lr()[0]))
train_gcn(model, criterion_xent, criterion_xent_person, criterion_pair,
criterion_htri_filter, criterion_htri, criterion_permutation,
optimizer, trainloader, use_gpu)
#train(model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu)
if args.stepsize > 0: scheduler.step()
if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (
epoch + 1) == args.max_epoch:
print("==> Test")
rank1 = test_gcn_person_batch(model, queryloader,
querygalleryloader, galleryloader,
args.pool, use_gpu)
#rank1 = test_gcn(model, queryloader, galleryloader, args.pool, use_gpu=False)
#rank1 = test(model, queryloader, galleryloader, args.pool, use_gpu)
is_best = rank1 > best_rank1
if is_best: best_rank1 = rank1
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'))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
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 main():
torch.manual_seed(args.seed) # 为CPU设置种子用于生成随机数，以使得结果是确定的
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices # 在代码中指定需要使用的GPU
use_gpu = torch.cuda.is_available() # 查看当前环境是否支持CUDA,支持返回true，不支持返回false
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: # 如果使用gpu，输出选定的gpu，
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True # 在程序刚开始加这条语句可以提升一点训练速度,没什么额外开销
torch.cuda.manual_seed_all(args.seed) # 为GPU设置种子用于生成随机数，以使得结果是确定的
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_dataset(name=args.dataset) # 初始化数据集,从data_manager.py文件中加载。
# import transforms as T.
# T.Compose=一起组合几个变换。
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width), # 以一个概率进行，首先将图像大小增加到（1 + 1/8），然后执行随机裁剪。
T.RandomHorizontalFlip(), # 以给定的概率(0.5)随机水平翻转给定的PIL图像。
T.ToTensor(), # 将``PIL Image``或``numpy.ndarray``转换为张量。
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # 用平均值和标准偏差归一化张量图像。
# input[channel] = (input[channel] - mean[channel]) / std[channel]
])
transform_test = T.Compose([
T.Resize((args.height, args.width)), # 将输入PIL图像的大小调整为给定大小。
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
# 设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。
pin_memory = True if use_gpu else False
# DataLoader数据加载器。 组合数据集和采样器，并在数据集上提供单进程或多进程迭代器。
trainloader = DataLoader(
# VideoDataset:基于视频的person reid的数据集.(训练的数据集，视频序列长度，采样方法：随机，进行数据增强）
VideoDataset(dataset.train, seq_len=args.seq_len, sample='random', transform=transform_train),
# 随机抽样N个身份，然后对于每个身份，随机抽样K个实例，因此批量大小为N * K.
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,
) # 如果数据集大小不能被批量大小整除，则设置为“True”以删除最后一个不完整的批次。
queryloader = DataLoader(
VideoDataset(dataset.query, seq_len=args.seq_len, sample='dense', transform=transform_test),
batch_size=args.test_batch,
shuffle=False, # 设置为“True”以使数据在每个时期重新洗牌（默认值：False）。
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False, # 如果“False”和数据集的大小不能被批量大小整除，那么最后一批将更小。
)
galleryloader = DataLoader(
VideoDataset(dataset.gallery, seq_len=args.seq_len, sample='dense', 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)) # 模型的初始化
if args.arch == 'resnet503d':
model = resnet3d.resnet50(num_classes=dataset.num_train_pids, sample_width=args.width,
sample_height=args.height, sample_duration=args.seq_len)
# 如果不存在预训练模型，则报错
if not os.path.exists(args.pretrained_model):
raise IOError("Can't find pretrained model: {}".format(args.pretrained_model))
# 导入预训练的模型
print("Loading checkpoint from '{}'".format(args.pretrained_model))
checkpoint = torch.load(args.pretrained_model)
state_dict = {} # 状态字典,从checkpoint文件中加载参数
for key in checkpoint['state_dict']:
if 'fc' in key:
continue
state_dict[key.partition("module.")[2]] = checkpoint['state_dict'][key]
model.load_state_dict(state_dict, strict=False)
else:
model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'})
print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters())/1000000.0))
# 损失函数：xent：softmax交叉熵损失函数。htri：三元组损失函数。
criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu)
criterion_htri = TripletLoss(margin=args.margin)
# 优化器：adam
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
# stepsize，逐步减少学习率（> 0表示已启用）
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma)
# lr_scheduler学习率计划，StepLR,将每个参数组的学习速率设置为每个步长时期由gamma衰减的初始lr.
start_epoch = args.start_epoch # 手动时期编号（重启时有用）
if use_gpu:
model = nn.DataParallel(model).cuda() # 多GPU训练
# DataParallel是torch.nn下的一个类，需要制定的参数是module（可以多gpu运行的类函数）和input（数据集）
if args.evaluate: # 这里的evaluate没有意义，应该添加代码导入保存的checkpoint，再test
print("Evaluate only") # 进行评估
test(model, queryloader, galleryloader, args.pool, use_gpu)
return
start_time = time.time() # 开始的时间
best_rank1 = -np.inf # 初始化，负无穷
if args.arch == 'resnet503d': # 如果模型为resnet503d,
torch.backends.cudnn.benchmark = False
for epoch in range(start_epoch, args.max_epoch): # epoch,从开始到最大，进行训练。
print("==> Epoch {}/{}".format(epoch+1, args.max_epoch))
train(model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu)
if args.stepsize > 0:
scheduler.step()
# 如果运行一次评估的需要的epoch数大于0，并且当前epoch+1能整除这个epoch数，或者等于最大epoch数。那么就进行一次评估。
if 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 # 比较，大于则返回true，否则返回false。
if is_best:
best_rank1 = rank1
if use_gpu:
state_dict = model.module.state_dict()
# 函数static_dict()用于返回包含模块所有状态的字典，包括参数和缓存。
else:
state_dict = model.state_dict()
# 保存checkpoint文件
save_checkpoint({
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
}, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch+1) + '.pth.tar'))
# 经过的时间
elapsed = round(time.time() - start_time) # round() 方法返回浮点数x的四舍五入值
elapsed = str(datetime.timedelta(seconds=elapsed)) # 对象代表两个时间之间的时间差,
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
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_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 attr_main():
stale_step = args.stalestep
runId = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
args.save_dir = os.path.join(args.save_dir, runId)
if not os.path.exists(args.save_dir):
os.mkdir(args.save_dir)
print(args.save_dir)
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('./log_train_' + runId + '.txt')
else:
sys.stdout = Logger('./log_test_' + runId + '.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_dataset(name=args.dataset, min_seq_len=args.seq_len, attr=True)
args.attr_lens = [[], [9, 10, 2]] #dataset.attr_lens
args.columns = ["upcolor", "downcolor", "gender"] #dataset.columns
print("Initializing model: {}".format(args.arch))
# if args.arch == "resnet50ta_attr" or args.arch == "resnet50ta_attr_newarch":
if args.arch == 'attr_resnet503d':
model = models.init_model(name=args.arch,
attr_lens=args.attr_lens,
model_type=args.model_type,
num_classes=dataset.num_train_pids,
sample_width=args.width,
sample_height=args.height,
sample_duration=args.seq_len)
torch.backends.cudnn.benchmark = False
else:
model = models.init_model(name=args.arch,
attr_lens=args.attr_lens,
model_type=args.model_type)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
if args.dataset == "duke":
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])
])
elif args.dataset == "mars":
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
# if args.predict:
# attr_predict(model, transform_test, use_gpu)
# return
# # For SQL mmir, just first_img_path/ random_img_path sampling works
# # For other mmir, call first_img/ random_img sampling
# mmirImgQueryLoader = VideoDataset(dataset.mmir_img_query, seq_len=1, # dataset.mmir_img_query[0:2]
# sample='random_img_path', # sample='random', #changing to random to get the imgs for running attr. extraction model, does the img_path match with 'random_img_path' sampling -- yap, the samples returned with both sampling matches in postgres and here, so next just get those images and run the attribute ex. models on them
# transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# # The above function basically just tells how we sample dfrom the img list of this specific tracklet (has pid, camid, attrs)
# # sampling more from gallery as test and gallery has same person, train has different persons
# # during retrieval same person with same features should have higher rank than different person with same features
# # 2/3 for image, rest is for video
# mmirImgQueryLoader = VideoDataset(
# dataset.mmir_img_gallery,
# seq_len=1, # dataset.mmir_img_query[0:2]
# sample='random_img_path', # sample='random',
# transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# print(len(dataset.mmir_img_gallery))
# print(len(dataset.mmir_video_gallery))
# # import psycopg2
# # # set-up a postgres connection
# # conn = psycopg2.connect(database='ng', user='******',password='******',
# # host='146.148.89.5', port=5432)
# # dbcur = conn.cursor()
# # print("connection successful")
# # for batch_idx, (pids, camids, attrs, img_path) in enumerate(tqdm(mmirImgQueryLoader)):
# # #print(pids)
# # images = list()
# # images.append(img_path)
# # # just save the img in img_path
# # sql = dbcur.mogrify("""
# # INSERT INTO mmir_ground (
# # mgid,
# # ctype,
# # pid,
# # camid,
# # ubc,
# # lbc,
# # gender,
# # imagepaths
# # ) VALUES (DEFAULT, %s,%s,%s,%s,%s,%s,%s) ON CONFLICT DO NOTHING;""", ("image",
# # str(pids),
# # str(camids),
# # str(attrs[0]),
# # str(attrs[1]),
# # str(attrs[2]), images)
# # )
# # # print(sql)
# # dbcur.execute(sql)
# # conn.commit()
# # print(dataset.mmir_video_query[0]) # if this works, for sdml or other mmir training will just take some part of train for img, some for video
# # For SQL mmir, just first_img_path sampling works, as not even using the image_array, using just the attributes
# # For other mmir, call random sampling
# # sampling the whole video has no impact in mmir, as we cannot take avg. of the separated lists attributes "accuracy"
# # we need frames of whose attributes would be one single one, and then we search in img and text that attributes
# mmirVideoQueryLoader = VideoDataset(dataset.mmir_video_query, seq_len=10, # 100 * 10 = 1000 frames as image
# sample='random_video', transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# mmirVideoQueryLoader = VideoDataset(dataset.mmir_video_gallery, seq_len=6, # 100 * 10 = 1000 frames as image
# sample='random_video', transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# # for batch_idx, (pids, camids, attrs, img_paths) in enumerate(tqdm(mmirVideoQueryLoader)):
# # #print(pids)
# # sql = dbcur.mogrify("""
# # INSERT INTO mmir_ground (
# # mgid,
# # ctype,
# # pid,
# # camid,
# # ubc,
# # lbc,
# # gender,
# # imagepaths
# # ) VALUES (DEFAULT, %s,%s,%s,%s,%s,%s,%s) ON CONFLICT DO NOTHING;""", ("video",
# # str(pids),
# # str(camids),
# # str(attrs[0]),
# # str(attrs[1]),
# # str(attrs[2]), img_paths)
# # )
# # dbcur.execute(sql)
# # conn.commit()
# # conn.close()
# # if i want to add more mmir image, video samples,
# # 1. in line 116, pass num_mmir_query_imgs=1000, num_mmir_query_videos=100
# # 2. just uncomment the SQL query code, the first 1000 & 100 would be same, as seed is set, the later ones will be added in postgres
# # hahahah, way surpassed above steps, created new samples from query_id, appended them; then created samples from gallery_id, added them by calling sampler again
# trainloader = DataLoader(
# VideoDataset(dataset.train, seq_len=args.seq_len, sample='random', transform=transform_train, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens),
# 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,
# )
# validloader = VideoDataset(dataset.valid, seq_len=args.seq_len,
# sample='dense', transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# #queryloader = VideoDataset(dataset.query + dataset.gallery, seq_len=args.seq_len, sample='single' transform=transform_test, attr=True,
# queryloader = VideoDataset(dataset.query + dataset.gallery, seq_len=args.seq_len, #args.seq_len, sample='dense'
# sample='dense', transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# # queryLoaderMIT = VideoDataset(dataset.query + dataset.gallery, seq_len=1,
# # sample='random_img_path', transform=transform_test, attr=True,
# # attr_loss=args.attr_loss, attr_lens=args.attr_lens)
start_epoch = args.start_epoch
if use_gpu:
model = nn.DataParallel(model).cuda()
start_time = time.time()
if args.arch == 'resnet503d':
torch.backends.cudnn.benchmark = False
# print("Run attribute pre-training")
if args.attr_loss == "cropy":
criterion = nn.CrossEntropyLoss()
elif args.attr_loss == "mse":
criterion = nn.MSELoss()
if args.colorsampling:
attr_test_MIT(model, criterion, queryLoaderMIT, use_gpu)
return
if args.evaluate:
print("Evaluate only")
# model_root = "/data/chenzy/models/mars/2019-02-26_21-02-13"
model_root = args.eval_model_root
model_paths = []
for m in os.listdir(model_root):
if m.endswith("pth"):
model_paths.append(m)
model_paths = sorted(model_paths,
key=lambda a: float(a.split("_")[1]),
reverse=True)
#print(model_paths)
#input()
# model_paths = ['rank1_2.8755379380596713_checkpoint_ep500.pth']
for m in model_paths[0:1]:
model_path = os.path.join(model_root, m)
print(model_path)
old_weights = torch.load(model_path)
new_weights = model.module.state_dict()
for k in new_weights:
if k in old_weights:
new_weights[k] = old_weights[k]
model.module.load_state_dict(new_weights)
if args.predict:
attr_predict(model, transform_test, use_gpu)
return
avr_acc = attr_test(model, criterion, queryloader, use_gpu)
# break
# test(model, queryloader, galleryloader, args.pool, use_gpu)
return
if use_gpu:
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.module.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
# avr_acc = attr_test(model, criterion, queryloader, use_gpu)
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=args.stepsize,
gamma=args.gamma)
best_avr = 0
no_rise = 0
for epoch in range(start_epoch, args.max_epoch):
print("==> Epoch {}/{}".format(epoch + 1, args.max_epoch))
attr_train(model, criterion, optimizer, trainloader, use_gpu)
if args.stepsize > 0: scheduler.step()
if args.eval_step > 0 and ((epoch + 1) % (args.eval_step) == 0 or
(epoch + 1) == args.max_epoch):
# avr_acc = attr_test(model, criterion, queryloader, use_gpu)
avr_acc = attr_test(model, criterion, validloader, use_gpu)
print("avr", avr_acc)
if avr_acc > best_avr:
no_rise = 0
print("==> Test")
best_avr = avr_acc
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
torch.save(
state_dict,
osp.join(
args.save_dir, "avr_" + str(avr_acc) +
'_checkpoint_ep' + str(epoch + 1) + '.pth'))
else:
no_rise += 1
print("no_rise:", no_rise)
if no_rise > stale_step:
break
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
"""
dataset = dataset_manager.init_img_dataset(
root='data',name=dataset_name,split_id=split_id,
)
"""
##### CUHK03
dataset = dataset_manager.init_img_dataset(
root='data',
name=dataset_name,
split_id=split_id,
cuhk03_labeled=cuhk03_labeled,
cuhk03_classic_split=cuhk03_classic_split,
)
tfms_train = tfms.Compose([
tfms.Random2DTranslation(256, 128),
tfms.RandomHorizontalFlip(),
tfms.ToTensor(),
tfms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
tfms_test = tfms.Compose([
tfms.Resize((256, 128)),
tfms.ToTensor(),
tfms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True
trainloader = DataLoader(
ImageDataset(dataset.train, transform=tfms_train),
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'))
# tensorboardX
# writer = SummaryWriter(log_dir=osp.join(args.save_dir,'summary'))
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_img_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]),
])
if args.random_erasing:
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]),
RandomErasing(probability=args.probability, mean=[0.0, 0.0, 0.0]),
])
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
if args.loss == 'xent,htri':
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train),
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,
)
elif args.loss == 'xent':
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train),
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),
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=args.loss)
print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters())/1000000.0))
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)
if args.stepsize > 0:
if not args.warmup:
scheduler = lr_scheduler.StepLR(optimizer, step_size=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, use_gpu)
return
def adjust_lr(optimizer, ep):
if ep < 20:
lr = 1e-4 * (ep + 1) / 2
elif ep < 80:
#lr = 1e-3 * len(args.gpu_devices)
lr = 1e-3
elif ep < 180:
#lr = 1e-4 * len(args.gpu_devices)
lr = 1e-4
elif ep < 300:
#lr = 1e-5 * len(args.gpu_devices)
lr = 1e-5
elif ep < 320:
#lr = 1e-5 * 0.1 ** ((ep - 320) / 80) * len(args.gpu_devices)
lr = 1e-5 * 0.1 ** ((ep - 320) / 80)
elif ep < 400:
lr = 1e-6
elif ep < 480:
#lr = 1e-4 * len(args.gpu_devices)
lr = 1e-4
else:
#lr = 1e-5 * len(args.gpu_devices)
lr = 1e-5
for p in optimizer.param_groups:
p['lr'] = lr
length = len(trainloader)
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
#best_rerank1 = -np.inf
#best_rerankepoch = 0
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
if args.stepsize > 0:
if args.warmup:
adjust_lr(optimizer, epoch + 1)
else:
scheduler.step()
train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu=use_gpu, summary=None, length=length)
train_time += round(time.time() - start_train_time)
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(epoch, model, queryloader, galleryloader, use_gpu=True, summary=None)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_epoch = epoch + 1
####### Best Rerank
#is_rerankbest = rerank1 > best_rerank1
#if is_rerankbest:
# best_rerank1 = rerank1
# best_rerankepoch = 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'))
writer.close()
print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(best_rank1, best_epoch))
#print("==> Best Rerank-1 {:.1%}, achieved at epoch {}".format(best_rerank1, best_rerankepoch))
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))
device = torch.device("cpu")
sys.stdout = Logger(osp.join(PATH, 'log_train.txt'))
print("Dataset is being initialized")
dataset = dataset_manager.init_img_dataset(
root='data',
name=dataset_name,
split_id=split_id,
cuhk03_labeled=cuhk03_labeled,
cuhk03_classic_split=cuhk03_classic_split,
)
tfms_train = tfms.Compose([
tfms.Random2DTranslation(height, width),
tfms.RandomHorizontalFlip(),
tfms.ToTensor(),
tfms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
tfms_test = tfms.Compose([
tfms.Resize(size=(height, width)),
tfms.ToTensor(),
tfms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True
trainloader = DataLoader(
ImageDataset(dataset.train, transform=tfms_train),
def main():
global Temporal_corr_prob
global paths, sorted_frame_list_all, max_frame
Temporal_corr_prob = np.zeros((9, 9, 5))
Temporal_corr_prob = np.load('./t_5_corr.npy')
Temporal_corr_prob[8, :, -1] = [
0.51415094, 0.20754717, 0.07075472, 0.25, 0.08962264, 0.45283019,
0.0754717, 0.49528302, 1
]
#Temporal_corr_prob[8,:,-1] = [0.96295517, 0.96295517, 0.64197011, 0.96295517, 0.77036413, 0.96295517,0.64197011, 0.96295517, 1]
#800 Temporal_corr_prob[8,:,-1] = [0.10991234, 0.03641268, 0.01213756, 0.04517869, 0.01416049, 0.09844909,0.01213756, 0.10249494,1]
#480 Temporal_corr_prob[8,:,-1] = [0.34993271 , 0.12651413 , 0.04306864, 0.16689098 , 0.05114401, 0.30417227, 0.04306864, 0.32570659,1]
#720 Temporal_corr_prob[8,:,-1] = [0.45967742, 0.1733871, 0.06451613, 0.22983871, 0.07258065, 0.40725806, 0.06451613, 0.43951613,1]
#600 Temporal_corr_prob[8,:,-1] = [0.4169468,0.15467384,0.0537996,0.19166106,0.0672495,0.35642233,0.05716207,0.39340955,1]
#1080 Temporal_corr_prob[8,:,-1] = [0.55533199, 0.23541247, 0.09054326, 0.28973843, 0.11468813, 0.52515091,0.09657948, 0.58551308,1]
#Temporal_corr_prob[8,:,-1] = [0.49708912, 0.17913121, 0.06717421, 0.2507837, 0.08060905, 0.43439319, 0.07165249, 0.46574116, 1]
#Temporal_corr_prob[8,2,29] += 0.1
#Temporal_corr_prob[8,8,29] -= 0.1
#Temporal_corr_prob[:,:,0:-1] = 0
#Temporal_corr_prob[:,:,-1] = Spatial_corr_prob
#Temporal_corr_prob[8,:,-1] = [0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1]
dataset = data_manager.init_img_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 = False
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train),
batch_size=args.train_batch,
shuffle=False,
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,
)
#gallery = ImageDatasetLazy(dataset.gallery, transform=transform_test)
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,
)
paths = dict()
life_frequency = []
stand_frequency = {0: [], 1: [], 2: [], 3: [], 4: [], 5: [], 6: [], 7: []}
frame_list = {
0: dict(),
1: dict(),
2: dict(),
3: dict(),
4: dict(),
5: dict(),
6: dict(),
7: dict()
}
change = np.zeros((9, 9))
cal_v = np.zeros(9)
total_number = 0
min_frame = 10000000
max_frame = -1
with torch.no_grad():
for batch_idx, (names, imgs, pids, camids,
fids) in enumerate(trainloader):
for s_index in range(len(names)):
if paths.get(int(pids[s_index])) == None:
paths[int(pids[s_index])] = []
fid_ = fids[s_index] + cam_offsets[camids[s_index]]
#if paths[int(pids[s_index])] != []:
# if paths
paths[int(pids[s_index])].append((fid_, camids[s_index]))
if int(fid_) >= max_frame:
max_frame = int(fid_)
if int(fid_) <= min_frame:
min_frame = int(fid_)
if frame_list[int(camids[s_index])].get(int(fid_)) == None:
frame_list[int(camids[s_index])][int(fid_)] = []
frame_list[int(camids[s_index])][int(fid_)].append(
int(pids[s_index]))
cal_v[camids[s_index]] += 1
total_number += 1
print("Max frame and min frame : ", max_frame, min_frame)
print("")
sorted_frame_list_all = [
] #{ 0 : list(), 1:list(),2:list(),3:list(),4:list(),5:list(),6:list(),7:list()}
for ind in range(8):
for key in list(frame_list[ind].keys()):
tmp_list = frame_list[ind][key]
tmp_list.insert(0, ind)
tmp_list.insert(0, key)
#sorted_frame_list[ind].append(tmp_list)
sorted_frame_list_all.append(tmp_list)
#sorted_frame_list[ind] = sorted(sorted_frame_list[ind], key=lambda x: x[0])
#print("Sorted Index : ",ind, "With term : ",len(sorted_frame_list[ind]))
sorted_frame_list_all = sorted(sorted_frame_list_all, key=lambda x: x[0])
#input()
print("Sorted Index : ", ind, "With term : ", len(sorted_frame_list_all))
#227540 49700
#version1(paths,sorted_frame_list)
#baseline(paths,sorted_frame_list)
'''global frame_window_size
global threshold_value
a = np.zeros(100)
b = np.zeros(100)
for i in range(100):
a[i] = 10*i + 10
b[i] = pow(0.42,i)
result = []
for i in range(100):
for j in range(100):
frame_window_size = a[i]
threshold_value = b[j]
t,f = version3(paths,sorted_frame_list_all,max_frame)
print("Time : ",i,j,"with tot : ",t, " and f : ",f, " setting : ",frame_window_size,threshold_value)
if f < 50 and t < 600000:
result.append((t,f))
print(result)'''
from bayes_opt import BayesianOptimization
# Bounded region of parameter space
#pbounds = {'t0':(0.6,1.5),'t1':(0.6,1.5),'t2':(0.2,0.8),'t3':(0.6,1.2),'t4':(0.2,0.8),'t5':(0.5,1.2),'t6':(0.2,0.8),'t7':(0.3,1.2),}
#wrapper =
for i in [
0.96, 0.965, 0.97, 0.975, 0.98, 0.99
]: #1.2,1.3,1.4,1.5,1.6,1.7,1.8,1.9,2.0,2.1,2.2,2.3,2.4,2.8,3,4,5,6,10,100,1e5,1e6]:#0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.735,0.80,0.85,0.9,
#for k in [7,13,17,20,23,27,31,71,103,143]:
j = i
t, f = wrapper(j)
print(j, ",", t, ",", f)
#version2(paths,sorted_frame_list_all,max_frame,j)
'''optimizer = BayesianOptimization(