def main():
global args
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu:
use_gpu = False
if use_gpu:
print('Currently using GPU {}'.format(args.gpu_devices))
cudnn.benchmark = True
else:
warnings.warn(
'Currently using CPU, however, GPU is highly recommended')
print('Initializing image data manager')
dm = ImageDataManager(use_gpu, **testset_kwargs(args)) # 创建 data manager
testloader_dict = dm.return_dataloaders(
) # 创建data loader字典（包含query和gallery）
print('Initializing model: {}'.format(
args.arch)) # args.arch默认是'resnet101'
model = models.init_model(
name=args.arch,
loss={'xent', 'htri'},
pretrained=not args.no_pretrained,
use_gpu=use_gpu) #models/__init__.py -> models/resnet.py/def resnet101
print('Model size: {:.3f} M'.format(count_num_param(model)))
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model,
args.load_weights) #下载好的model.pth.tar-9文件
model = nn.DataParallel(model).cuda() if use_gpu else model #多GPU
print('Matching {} ...'.format(args.test_set)) #split好的query和gallery
queryloader = testloader_dict['query']
galleryloader = testloader_dict['test']
distmat, q_pids, g_pids, q_camids, g_camids = run(
model, queryloader, galleryloader, use_gpu,
return_distmat=True) #RE-ID
np.savez(args.save_npy,
distmat=distmat,
q_pids=q_pids,
g_pids=g_pids,
q_camids=q_camids,
g_camids=g_camids) #存储结果到distmatrix.npz后续可以本地跑出track3.txt
# 到这为止，利用定义的reid模型，读取extract的img，得到距离矩阵distmat和 行所代表的queryID+camID 和列所代表的galleryID+camID
#下一步，后处理，解析距离矩阵的物理含义，得到MTMC的结果
result = np.load(args.save_npy) #存储完马上加载distmatrix.npz
reid_dict, rm_dict = calc_reid(
result) #得到reid可能匹配结果reid_dict和reid无匹配结果rm_dict
print(rm_dict, reid_dict)
with open(args.tracklet_path, "r") as f: #读取tracklets.txt
or_tracks = f.readlines()
g = open(args.track3_path, "w") #打开track3.txt
update_output(or_tracks, reid_dict, rm_dict,
g) #更新tracklets.txt->track3.txt
def main():
global args
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu:
use_gpu = False
if use_gpu:
print('Currently using GPU {}'.format(args.gpu_devices))
cudnn.benchmark = True
else:
warnings.warn(
'Currently using CPU, however, GPU is highly recommended')
print('Initializing image data manager')
dm = ImageDataManager(use_gpu, **testset_kwargs(args))
testloader_dict = dm.return_dataloaders()
print('Initializing model: {}'.format(args.arch))
model = models.init_model(name=args.arch,
loss={'xent', 'htri'},
pretrained=not args.no_pretrained,
use_gpu=use_gpu)
print('Model size: {:.3f} M'.format(count_num_param(model)))
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model, args.load_weights)
model = nn.DataParallel(model).cuda() if use_gpu else model
print('Matching {} ...'.format(args.test_set))
queryloader = testloader_dict['query']
galleryloader = testloader_dict['test']
distmat, q_pids, g_pids, q_camids, g_camids = run(model,
queryloader,
galleryloader,
use_gpu,
return_distmat=True)
np.savez(args.save_npy,
distmat=distmat,
q_pids=q_pids,
g_pids=g_pids,
q_camids=q_camids,
g_camids=g_camids)
result = np.load(args.save_npy)
reid_dict, rm_dict = calc_reid(result)
print(rm_dict, reid_dict)
with open(args.tracklet_path, "r") as f:
or_tracks = f.readlines()
g = open(args.track3_path, "w")
update_output(or_tracks, reid_dict, rm_dict, g)
def main():
# 为了看看repo提供的model.pth.tar-9在validation集的mAp和rank-1
# 我自己训练的tar-9只有mAP: 15.1%； Rank-1: 23.3% ，不知道为什么
# 更改args.load_weights = '/model/caohw9/track3_model/model.pth.tar-9'
global args
print(args)
set_random_seed(args.seed)
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu:
use_gpu = False
sys.stdout = Logger(osp.join(args.save_dir, "log.txt"))
if use_gpu:
print('Currently using GPU {}'.format(args.gpu_devices))
cudnn.benchmark = True
else:
warnings.warn(
'Currently using CPU, however, GPU is highly recommended')
# 初始化loader
print('Initializing image data manager')
dm = ImageDataManager(use_gpu, **trainset_kwargs(args))
trainloader, testloader_dict = dm.return_dataloaders(
) #trainloader用于训练，testloader_dict包含['query']和['gallery']2个loader
print('suffessfully initialize loaders!')
# 初始化模型
print('Initializing model: {}'.format(
args.arch)) #args.arch default='resnet101'
model = models.init_model(name=args.arch,
num_classes=dm.num_train_pids,
loss={'xent', 'htri'},
pretrained=not args.no_pretrained,
use_gpu=use_gpu)
print('Model size: {:.3f} M'.format(count_num_param(model)))
# 加载预训练参数
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model, args.load_weights)
#加载训练过的模型后，先看看validation
print('=> Validation')
print('Evaluating {} ...'.format(
args.test_set)) #args.test_set应该是指的validation set？
queryloader = testloader_dict['query']
galleryloader = testloader_dict['test']
model = nn.DataParallel(model).cuda() if use_gpu else model
rank1 = test(model, queryloader, galleryloader, use_gpu) #validation！
# 多GPU训练
else:
model = nn.DataParallel(model).cuda() if use_gpu else model
# 定义loss,optimizer, lr_scheduler
criterion_xent = CrossEntropyLoss(num_classes=dm.num_train_pids,
use_gpu=use_gpu,
label_smooth=args.label_smooth)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optimizer(model, **optimizer_kwargs(args))
scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs(args))
# 是否是resume训练
if args.resume and check_isfile(args.resume):
args.start_epoch = resume_from_checkpoint(
args.resume, model, optimizer=optimizer) #获取中断时刻的epoch数
# 开始训练！
time_start = time.time()
print('=> Start training')
for epoch in range(args.start_epoch, args.max_epoch):
train(epoch, model, criterion_xent, criterion_htri, optimizer,
trainloader, use_gpu) #训练
scheduler.step() #更新lr
# 当epoch数超过args.start_eval，每隔一定频率args.eval_freq，或者达到最后一个epoch，进行validation+存储checkpoint
if (epoch + 1) > args.start_eval and args.eval_freq > 0 and (
epoch + 1) % args.eval_freq == 0 or (epoch +
1) == args.max_epoch:
print('=> Validation')
print('Evaluating {} ...'.format(
args.test_set)) #args.test_set应该是指的validation set？
queryloader = testloader_dict['query']
galleryloader = testloader_dict['test']
rank1 = test(model, queryloader, galleryloader,
use_gpu) #validation！
save_checkpoint(
{
'state_dict': model.state_dict(), #模型的状态字典
'rank1': rank1,
'epoch': epoch + 1,
'arch': args.arch, #default='resnet101'
'optimizer': optimizer.state_dict(
), #优化器对象的状态字典，包含优化器的状态和超参数(如lr, momentum,weight_decay等)
},
args.save_dir) #validation同时保存checkpoint
# 训练结束！
elapsed = round(time.time() - time_start) #持续时间
elapsed = str(datetime.timedelta(seconds=elapsed))
print('Elapsed {}'.format(elapsed))
warnings.warn(
'Currently using CPU, however, GPU is highly recommended')
print('Initializing image data manager')
dm = ImageDataManager(use_gpu, **testset_kwargs(args))
testloader_dict = dm.return_dataloaders()
print('Initializing model: {}'.format(args.arch))
model = models.init_model(name=args.arch,
loss={'xent', 'htri'},
pretrained=not args.no_pretrained,
use_gpu=use_gpu)
print('Model size: {:.3f} M'.format(count_num_param(model)))
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model, args.load_weights)
model = nn.DataParallel(model).cuda() if use_gpu else model
print('Matching {} ...'.format(args.test_set))
queryloader = testloader_dict['query']
galleryloader = testloader_dict['test']
# distmat, q_pids, g_pids, q_camids, g_camids = run(model, queryloader, galleryloader, use_gpu, return_distmat=True)
batch_time = AverageMeter()
model.eval()
with torch.no_grad():
gf, g_pids, g_camids = [], [], []
for batch_idx, (imgs, pids, camids,
_) in enumerate(tqdm(galleryloader)):
def main():
global args
set_random_seed(args.seed)
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu:
use_gpu = False
sys.stdout = Logger(osp.join(args.save_dir, "log.txt"))
if use_gpu:
print('Currently using GPU {}'.format(args.gpu_devices))
cudnn.benchmark = True
else:
warnings.warn(
'Currently using CPU, however, GPU is highly recommended')
print('Initializing image data manager')
dm = ImageDataManager(use_gpu, **trainset_kwargs(args))
trainloader, testloader_dict = dm.return_dataloaders()
print('Initializing model: {}'.format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dm.num_train_pids,
loss={'xent', 'htri'},
pretrained=not args.no_pretrained,
use_gpu=use_gpu)
print('Model size: {:.3f} M'.format(count_num_param(model)))
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model, args.load_weights)
model = nn.DataParallel(model).cuda() if use_gpu else model
criterion_xent = CrossEntropyLoss(num_classes=dm.num_train_pids,
use_gpu=use_gpu,
label_smooth=args.label_smooth)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optimizer(model, **optimizer_kwargs(args))
scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs(args))
if args.resume and check_isfile(args.resume):
args.start_epoch = resume_from_checkpoint(args.resume,
model,
optimizer=optimizer)
time_start = time.time()
print('=> Start training')
for epoch in range(args.start_epoch, args.max_epoch):
train(epoch, model, criterion_xent, criterion_htri, optimizer,
trainloader, use_gpu)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_freq > 0 and (
epoch + 1) % args.eval_freq == 0 or (epoch +
1) == args.max_epoch:
print('=> Validation')
print('Evaluating {} ...'.format(args.test_set))
queryloader = testloader_dict['query']
galleryloader = testloader_dict['test']
rank1 = test(model, queryloader, galleryloader, use_gpu)
save_checkpoint(
{
'state_dict': model.state_dict(),
'rank1': rank1,
'epoch': epoch + 1,
'arch': args.arch,
'optimizer': optimizer.state_dict(),
}, args.save_dir)
elapsed = round(time.time() - time_start)
elapsed = str(datetime.timedelta(seconds=elapsed))
print('Elapsed {}'.format(elapsed))