def extract_feature(model, data_loader, use_gpu, print_freq=1, metric=None):
'''
:param model: the model has been trained
:param data_loader: query_loader / gallery_loader
:param print_freq: default 1
:param use_gpu: delivered from main.py
:param metric:
:return: feature & pids & camids
'''
model.eval()
batch_time = AverageMeter()
features, pids, camids = [], [], []
for batch_idx, (imgs, pids, camids) in enumerate(data_loader):
if use_gpu:
imgs = imgs.cuda()
end = time.time()
features = model(imgs) # resnet50: f
batch_time.update(time.time() - end)
features = features.data.cpu()
features.append(features)
pids.extend(pids)
camids.extend(camids)
features = torch.cat(features, 0)
pids = np.asarray(pids)
camids = np.asarray(camids)
print("Extracted features for {} set, obtained {}-by-{} matrix".format(
data_loader, features.size(0), features.size(1)))
return features, pids, camids
def train(train_loader, model, criterion, optimizer, epoch):
global opt
losses = AverageMeter()
# switch to train mode
model.train()
criterion.train()
step = epoch * len(train_loader)
pred_gt_same = []
for i, (box, cls, feature, lfeat, lrel, sents, sents_gt, gt_boxes, img_ids,
sent_ids) in enumerate(train_loader):
step += 1
if opt['gpus'] is not None:
box = box.cuda()
cls = cls.cuda()
feature = feature.cuda()
lfeat = lfeat.cuda()
lrel = lrel.cuda()
sents = sents.cuda()
sents_gt = sents_gt.cuda()
# compute output
score = model(feature, cls, lfeat, lrel, sents)
loss, score = criterion(score, box, cls, sents_gt)
losses.update(loss.item())
cls = to_numpy(cls)
final_score = to_numpy(score.detach())
final_score[cls == -1] = -999
pred_ind = np.argmax(final_score, 1)
sents_gt = to_numpy(sents_gt)
for j in range(pred_ind.size):
if sents_gt[j] == pred_ind[j]:
pred_gt_same.append(1)
else:
pred_gt_same.append(0)
# compute gradient and do Adam step
optimizer.zero_grad()
loss.backward()
clip_gradient(optimizer, opt['grad_clip'])
optimizer.step()
if i % args.print_freq == 0:
if i != 0:
same = np.sum(
pred_gt_same[-args.print_freq *
opt['batch_size']:]) / float(
args.print_freq * opt['batch_size'])
print('Epoch: [{0}][{1}/{2}]\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec {same:.4f}'.format(epoch,
i,
len(train_loader),
loss=losses,
same=same))
def train_model_for_epoch(model, train_dataloader, loss_fn, optimizer, epoch):
"""
train model for a epoch
:param model:
:param train_dataloader:
:param loss_fn:
:param optimizer:
:return:
"""
# prepared
model.train()
num_iter = len(train_dataloader) if cfg.max_iter <= 0 else min(len(train_dataloader), cfg.max_iter)
loss_states = loss_fn.get_loss_states()
data_time, batch_time = AverageMeter(), AverageMeter()
running_loss = AverageMeter()
avg_loss_states = {l: AverageMeter() for l in loss_states}
start = time.time()
last = time.time()
phase = "train"
# foreach the images
for iter_id, train_data in enumerate(train_dataloader):
if iter_id >= num_iter:
break
# load data time
data_time.update(time.time() - last)
inputs, targets, infos = train_data
# to device
inputs = inputs.to(device)
# forward the models and loss
outputs = model(inputs)
loss, loss_stats = loss_fn(outputs, targets)
# update the weights
optimizer.zero_grad()
loss.backward()
optimizer.step()
# network time and update time
batch_time.update(time.time() - last)
last = time.time()
# handle the log and accumulate the loss
# logger.open_summary_writer()
log_item = '{phase} per epoch: [{0}][{1}/{2}]|Tot: {total:} '.format(
epoch, iter_id, num_iter, phase=phase, total=last - start)
for l in avg_loss_states:
if l in loss_stats:
avg_loss_states[l].update(
loss_stats[l].item(), inputs.size(0))
log_item = log_item + '|{}:{:.4f}'.format(l, avg_loss_states[l].avg)
# logger.scalar_summary('{phase}/epoch/{}'.format(l, phase=phase), avg_loss_states[l].avg, epoch* num_iter + iter_id)
# logger.close_summary_writer()
running_loss.update(loss.item(), inputs.size(0))
log_item = log_item + '|Data {dt.val:.3f}s({dt.avg:.3f}s) ' \
'|Net {bt.avg:.3f}s'.format(dt=data_time, bt=batch_time)
logger.write(log_item, level=1)
del inputs, loss
torch.cuda.empty_cache()
if (iter_id + 1) % cfg.save_span == 0:
executor.submit(save_checkpoint, model.state_dict(), epoch, running_loss.avg, data_cfg.save_dir, iter_id)
return running_loss, avg_loss_states
def validate(val_loader, model, criterion, epoch=-1):
global opt
losses = AverageMeter()
# switch to eval mode
model.eval()
criterion.eval()
pred_gt_same = []
with torch.no_grad():
for i, (box, cls, feature, lfeat, lrel, sents, sents_gt, gt_boxes,
img_ids, sent_ids) in enumerate(val_loader):
if opt['gpus'] is not None:
box = box.cuda()
cls = cls.cuda()
feature = feature.cuda()
lfeat = lfeat.cuda()
lrel = lrel.cuda()
sents = sents.cuda()
sents_gt = sents_gt.cuda()
# compute output
score = model(feature, cls, lfeat, lrel, sents)
loss, score = criterion(score, box, cls, sents_gt)
losses.update(loss.item())
cls = to_numpy(cls)
final_score = to_numpy(score.detach())
final_score[cls == -1] = -999
pred_ind = np.argmax(final_score, 1)
sents_gt = to_numpy(sents_gt)
for j in range(pred_ind.size):
if sents_gt[j] == pred_ind[j]:
pred_gt_same.append(1)
else:
pred_gt_same.append(0)
if i % args.print_freq == 0:
if i != 0:
same = np.sum(
pred_gt_same[-args.print_freq *
opt['batch_size']:]) / float(
args.print_freq * opt['batch_size'])
print('Epoch: [{0}][{1}/{2}]\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec {same:.4f}'.format(epoch,
i,
len(val_loader),
loss=losses,
same=same))
same = np.sum(pred_gt_same) / float(len(pred_gt_same))
print('Epoch: [{0}]\t'
'Loss {1:.4f}\t'
'Prec {2:.4f}'.format(epoch, losses.avg, same))
return losses.avg, same
def evaluator(self, query_loader, gallery_loader, use_gpu, ranks=None):
'''
:param model: trained model
:param query_loader: query-larder
:param gallery_loader: gallery-loader
:param use_gpu: given from main.py
:param ranks: the ranks we need
:return: CMC / mAP
'''
self.model.eval()
batch_time = AverageMeter()
with torch.no_grad():
# get the feature / pids / camids from the data_loader
qf, q_pids, q_camids, batch_time = self._extract_feature(
query_loader, use_gpu, batch_time)
gf, g_pids, g_camids, batch_time = self._extract_feature(
gallery_loader, use_gpu, batch_time)
print("==> BatchTime(s)/BatchSize(img): {:.3f}/{}".format(
batch_time.avg, self.test_batch))
distmat = self._pair_dist(qf, gf)
print("==> Computing CMC & mAP & mINP <==")
def train(self,
train_dataloader: DataLoader,
val_dataloader: DataLoader,
epoches=100):
best_score = 0
early_n = 0
for epo in range(epoches):
step_n = 0
train_avg_loss = AverageMeter()
train_data_iter = tqdm.tqdm(train_dataloader)
for batch in train_data_iter:
self.model.train()
batch_data = self._batch_trans(batch)
train_loss = self.step(step_n, batch_data)
train_avg_loss.update(train_loss.item(), 1)
status = '[{0}] lr = {1:.7f} batch_loss = {2:.3f} avg_loss = {3:.3f} '.format(
epo + 1,
self.scheduler.get_lr()[0], train_loss.item(),
train_avg_loss.avg)
#if step_n%self.log_steps ==0:
# print(status)
train_data_iter.set_description(status)
step_n += 1
if self.global_step % self.val_steps == 0:
## val
m = self.val(val_dataloader)
acc = m['acc']
if best_score < acc:
early_n = 0
best_score = acc
model_path = os.path.join(self.save_dir, 'best.pth')
torch.save(self.model.state_dict(), model_path)
else:
early_n += 1
self.logger.write("steps: {} ,mean ap : {:.4f} , best ap: {:.4f}". \
format(self.global_step, acc, best_score))
self.logger.write(str(m))
self.logger.write("==" * 50)
if early_n > self.early_stop_n:
return best_score
return best_score
def train(self, train_loader, val_loader, optimizer, lr_scheduler,
tb_logger):
print_freq = self.config.logging.print_freq
flops = calc_adaptive_model_flops(self.model,
self.config.dataset.input_size)
params = calc_model_parameters(self.model)
self._info('flops: {}, params: {}'.format(flops, params))
# meters
batch_time = AverageMeter(print_freq)
data_time = AverageMeter(print_freq)
loss_meter = AverageMeter(print_freq)
top1_meter = AverageMeter(print_freq)
top5_meter = AverageMeter(print_freq)
meters = [top1_meter, top5_meter, loss_meter, data_time]
criterion = self._get_criterion()
end = time.time()
for e in range(self.cur_epoch, self.config.training.epoch):
# train
if self.config.distributed.enable:
train_loader.sampler.set_epoch(e)
for batch_idx, (x, y) in enumerate(train_loader):
self.model.train()
x, y = x.cuda(), y.cuda()
self._train_one_batch(x, y, optimizer, lr_scheduler, meters,
[criterion], end)
batch_time.update(time.time() - end)
end = time.time()
cur_lr = lr_scheduler.get_lr()[0]
self._logging(tb_logger, e, batch_idx, len(train_loader),
meters + [batch_time], cur_lr)
# validation
if self.cur_step >= self.config.validation.start_val and self.cur_step % self.config.validation.val_freq == 0:
best_top1 = self.best_top1
self.validate(val_loader, tb_logger=tb_logger)
save_file = self.save(optimizer,
e,
best_top1=self.best_top1)
if self.best_top1 > best_top1:
from shutil import copyfile
best_file_dir = os.path.join(self.config.save_path,
'best')
if not os.path.exists(best_file_dir):
os.makedirs(best_file_dir)
best_file = os.path.join(best_file_dir, 'best.pth')
copyfile(save_file, best_file)
def main(args):
# SYN Dataset
dataset = SYN(
root=args.root,
w=args.width,
h=args.height,
t=args.time,
dataset='train',
train=True,
avi_dir=args.avi_dir,
usual_transform=False,
)
# Pytorch dataloader
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=args.cuda,
collate_fn=my_collate)
# Loop
data_time = AverageMeter()
start_data = time.time()
for i, dict_input in enumerate(dataloader):
duration_data = time.time() - start_data
data_time.update(duration_data)
# Get the data
clip, skeleton = dict_input['clip'], dict_input[
'skeleton'] # (B, C, T, 224, 224), (B, T, 2, 25, 2)
# Show
show_one_img(clip[0, :, 0], skeleton[0, 0])
print("{}/{} : {time.val:.3f} ({time.avg:.3f}) sec/batch".format(
i + 1, len(dataloader), time=data_time))
sys.stdout.flush()
start_data = time.time()
def evaluate(conf, model, test_loader):
model.eval()
top1 = AverageMeter()
top5 = AverageMeter()
with torch.no_grad():
for i, data in enumerate(test_loader):
timages = data[0].type(torch.FloatTensor)
tlabels = data[1].type(torch.LongTensor)
if conf.get()['cuda']['avail']:
timages, tlabels = timages.to(
conf.get()['cuda']['device']), tlabels.to(
conf.get()['cuda']['device'])
toutputs = model(timages)
ttop1, ttop5 = accuracy(toutputs, tlabels, (1, 5))
top1.update(ttop1.item(), timages.size(0))
top5.update(ttop5.item(), timages.size(0))
print('<<<TEST>>> top1({:.4f}) top5({:.4f})'.format(
top1.avg, top5.avg))
def train(self, train_loader, val_loader, optimizer, lr_scheduler,
tb_logger):
print_freq = self.config.logging.print_freq
# meters
batch_time = AverageMeter(print_freq)
data_time = AverageMeter(print_freq)
# track stats of min width, max width and random width
loss_meter = [AverageMeter(print_freq) for _ in range(3)]
top1_meter = [AverageMeter(print_freq) for _ in range(3)]
top5_meter = [AverageMeter(print_freq) for _ in range(3)]
meters = [top1_meter, top5_meter, loss_meter, data_time]
criterions = self._get_criterion()
self._sample_width()
end = time.time()
for e in range(self.cur_epoch, self.config.training.epoch):
# train
self.model.train()
if self.config.distributed.enable:
train_loader.sampler.set_epoch(e)
for batch_idx, (x, y) in enumerate(train_loader):
x, y = x.cuda(), y.cuda()
self._train_one_batch(x, y, optimizer, lr_scheduler, meters,
criterions, end)
batch_time.update(time.time() - end)
end = time.time()
cur_lr = lr_scheduler.get_lr()[0]
self._logging(tb_logger, e, batch_idx, len(train_loader),
meters + [batch_time], cur_lr)
# validation
self.validate(val_loader,
train_loader,
self.config.validation.width,
tb_logger=tb_logger)
self.save(optimizer, lr_scheduler, e)
def evaluate(dataloader, model, dev, topk=(1, )):
"""
:param dataloader:
:param model:
:param dev: devices, gpu or cpu
:param topk: [tuple] output the top topk accuracy
:return: [list[float]] topk accuracy
"""
model.eval()
test_accuracy = AverageMeter()
test_accuracy.reset()
with torch.no_grad():
for _, sample in enumerate(tqdm(dataloader, ncols=100, ascii=' >')):
x = sample['data'].to(dev)
y = sample['label'].to(dev)
output = model(x)
logits = output['logits']
acc = accuracy(logits, y, topk)
test_accuracy.update(acc[0], x.size(0))
return test_accuracy.avg
def do_train(cfg, model, center_criterion, train_loader, val_loader, optimizer,
optimizer_center, scheduler, loss_fn, num_query, last_epoch):
log_period = cfg.SOLVER.LOG_PERIOD
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
device = "cuda"
epochs = cfg.SOLVER.MAX_EPOCHS
logger = logging.getLogger("reid_baseline.train")
logger.info('start training')
if device:
model.to(device)
if torch.cuda.device_count() > 1:
print('Using {} GPUs for training'.format(
torch.cuda.device_count()))
model = nn.DataParallel(model)
else:
if cfg.SOLVER.FP16:
model, optimizer = amp.initialize(model,
optimizer,
opt_level='O1')
loss_meter = AverageMeter()
acc_meter = AverageMeter()
# train
for epoch in range(last_epoch, epochs + 1):
start_time = time.time()
loss_meter.reset()
acc_meter.reset()
model.train()
try:
for n_iter, (img, vid) in enumerate(train_loader):
optimizer.zero_grad()
optimizer_center.zero_grad()
img = img.to(device)
target = vid.to(device)
score, feat = model(img, target)
loss = loss_fn(score, feat, target)
if cfg.SOLVER.FP16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
if 'center' in cfg.MODEL.METRIC_LOSS_TYPE:
for param in center_criterion.parameters():
param.grad.data *= (1. / cfg.SOLVER.CENTER_LOSS_WEIGHT)
optimizer_center.step()
acc = (score.max(1)[1] == target).float().mean()
loss_meter.update(loss.item(), img.shape[0])
acc_meter.update(acc, 1)
if (n_iter + 1) % log_period == 0:
logger.info(
"Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Acc: {:.3f}, Base Lr: {:.2e}"
.format(epoch, (n_iter + 1), len(train_loader),
loss_meter.avg, acc_meter.avg,
scheduler.get_lr()[0]))
scheduler.step()
end_time = time.time()
time_per_batch = (end_time - start_time) / (n_iter + 1)
logger.info(
"Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]"
.format(epoch, time_per_batch,
train_loader.batch_size / time_per_batch))
if epoch % checkpoint_period == 0:
torch.save(
{
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict()
},
os.path.join(cfg.OUTPUT_DIR,
cfg.MODEL.NAME + '_{}.pth'.format(epoch)))
except:
torch.save(
{
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict()
},
os.path.join(cfg.OUTPUT_DIR,
cfg.MODEL.NAME + '_{}.pth'.format(epoch)))
if args.model == 'Q':
model = Qmodel(Ncls)
elif args.model == 'I':
model = Imodel(Ncls)
elif args.model == 'QI':
model = QImodel(Ncls)
elif args.model == 'RN':
model = RN(Ncls)
model.type(dtype)
print(model)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
Nepochs = args.Nepochs
Modelsavefreq = args.savefreq
tracker = AverageMeter()
early_stop = EarlyStopping(monitor='loss', min_delta=0, patience=5)
saveloss = defaultdict(list)
for epoch in range(0, Nepochs):
train_kwargs = {
'net': model,
'split': 'train',
'loader': trainloader,
'optimizer': optimizer,
'tracker': tracker,
'epoch': epoch
}
test_kwargs = {
def train_quantize(runner):
runner.reset_average_meters()
ckpt = runner.load_checkpoint("train", "best")
runner.model.load_state_dict(ckpt)
runner.model.fuse_model()
runner.model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
torch.quantization.prepare_qat(runner.model, inplace=True)
for epoch in range(runner.quantize_epochs):
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
#train
runner.model.train()
for i, data in enumerate(tqdm(runner.tr_loader)):
images, labels = data
if conf.get()['cuda']['avail']:
images, labels = images.to(runner.device), labels.to(
torch.device(runner.device))
runner.model = runner.model.to(torch.device(runner.device))
runner.optimizer.zero_grad()
outputs, loss = runner.regularizer(images, labels)
loss.backward()
runner.optimizer.step()
ttop1, ttop5 = runner.accuracy(outputs, labels, (1, 5))
losses.update(loss.item(), images.size(0))
top1.update(ttop1.item(), images.size(0))
top5.update(ttop5.item(), images.size(0))
print('[{:d}/{:d}] <<<TRAIN>>> lr({:.10f}) loss({:.4f}) top1({:.3f}) top5({:.3f})'.format( \
epoch + 1, runner.quantize_epochs, runner.optimizer.param_groups[0]['lr'], losses.avg, top1.avg, top5.avg) )
if epoch > 3:
# Freeze quantizer parameters
runner.model.apply(torch.quantization.disable_observer)
if epoch > 2:
# Freeze batch norm mean and variance estimates
runner.model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)
#evaluate
quantized_model = torch.quantization.convert(runner.model.eval(),
inplace=False)
quantized_model.eval()
etop1 = AverageMeter()
etop5 = AverageMeter()
with torch.no_grad():
for i, data in enumerate(
tqdm(
itertools.islice(runner.test_loader,
runner.quantize_evals))):
if conf.get()['data']['dali']['avail']:
timages = data[0]["data"]
tlabels = data[0]["label"].squeeze().long()
else:
timages = data[0].type(torch.FloatTensor)
tlabels = data[1].type(torch.LongTensor)
if runner.conf.get()['cuda']['avail']:
timages, tlabels = timages.to(runner.device), tlabels.to(
runner.device)
toutputs = runner.model(timages)
ttop1, ttop5 = runner.accuracy(toutputs, tlabels, (1, 5))
etop1.update(ttop1.item(), timages.size(0))
etop5.update(ttop5.item(), timages.size(0))
if runner.ebest_acc_top1 <= etop1.avg:
runner.ebest_acc_top1 = etop1.avg
if runner.ebest_acc_top5 < etop5.avg:
runner.ebest_acc_top5 = etop5.avg
model_path = runner.conf.get(
)['model']['path'] + "/" + runner.conf.get()['model']['name']
model_path += runner.fn_checkpoint("quantize", "best")
torch.jit.save(torch.jit.script(quantized_model), model_path)
print(
'[{:d}/{:d}] <<<TEST>>> top1({:.4f}) top5({:.4f}) best-top1({:.4f}) best-top5({:.4f})'
.format(epoch + 1, runner.quantize_epochs, etop1.avg,
etop5.avg, runner.ebest_acc_top1,
runner.ebest_acc_top5))
print('train_quantize() Done.')
def train_prune(runner):
model_prune_rate, weight_incremental = get_weight_fractions_with_target_ratio(prune_rate=runner.prune_rate, \
pruning_iterations=runner.prune_iterations, \
pruning_target_ratio=runner.prune_target_ratio)
print('Pruning iterations :', runner.prune_iterations, ', factions : ',
model_prune_rate)
runner.reset_average_meters()
for prune in range(runner.prune_iterations):
if prune == 0:
# load "train" model and convert it into "train_pruned" model
model_path = conf.get()['model']['path'] + "/" + conf.get(
)['model']['name']
checkpoint_file = model_path + "-tr-best.pt"
print("Load baseline checkpoint:", checkpoint_file)
checkpoint = torch.load(checkpoint_file)
runner.model.load_state_dict(checkpoint['model_state_dict'])
runner.micro_score()
else:
checkpoint_file = conf.get()['model']['path'] + "/" + conf.get(
)['model']['name'] + "-pr-" + str(prune - 1) + ".pt"
print("Load baseline checkpoint:", checkpoint_file)
ckpt = runner.load_checkpoint("train_prune", str(prune - 1))
runner.model.load_state_dict(ckpt)
runner.micro_score()
for epoch in range(runner.prune_epochs):
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
runner.model.train()
if epoch == 0:
masks = weight_prune(runner.model, model_prune_rate[prune])
runner.model.set_masks(masks)
print("Prune set_masks() : ", prune, "th iterations")
for i, data in enumerate(tqdm(runner.tr_loader)):
images, labels = data
if conf.get()['cuda']['avail']:
images, labels = images.to(runner.device), labels.to(
torch.device(runner.device))
runner.model = runner.model.to(torch.device(runner.device))
runner.optimizer.zero_grad()
outputs, loss = runner.regularizer(images, labels)
loss.backward()
runner.optimizer.step()
ttop1, ttop5 = runner.accuracy(outputs, labels, (1, 5))
losses.update(loss.item(), images.size(0))
top1.update(ttop1.item(), images.size(0))
top5.update(ttop5.item(), images.size(0))
print(
'[{:d}/{:d}] lr({:.10f}) loss({:.4f}) top1({:.3f}) top5({:.3f})'
.format(epoch + 1, runner.prune_epochs,
runner.optimizer.param_groups[0]['lr'], losses.avg,
top1.avg, top5.avg))
runner.evaluate("train_prune", epoch)
runner.scheduler.step()
runner.save_checkpoint("train_prune", str(prune))
if conf.get()['model']['pruning']['type'] == 'iterative_prune':
runner.tr_loader, _, runner.test_loader = get_dataloader(conf)
# runner.model.remove_buffers() # Remove temporary registered buffers(e.g., MaskedLinear, MaskedConv2d)
print('train_prune() Done.')
def train(trn_loader, model, criterion, optimizer, scheduler, epoch, G, args):
"""Train the classifier for a single epoch.
Params:
trn_loader: Loader for training set.
model: Classifier instance.
criterion: Loss function.
optimizer: Optimization algorithm to use for training.
scheduler: Learning rate scheduler.
epoch: Current training epoch.
G: Model of natural variation.
args: Command line arguments.
"""
net_meter = NetworkMeter()
timer = TimeMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
model.train()
for i, (imgs, target) in enumerate(trn_loader):
if args.short_epoch and (i > 10): break
imgs, target = imgs.cuda(), target.cuda()
batch_num = i + 1
timer.batch_start()
scheduler.update_lr(epoch, i + 1, len(trn_loader))
if args.mda is True:
imgs, target = mda_train(imgs, target, model, G, args)
elif args.mrt is True:
imgs, target = mrt_train(imgs, target, model, criterion, G, args)
elif args.mat is True:
imgs, target = mat_train(imgs, target, model, criterion, G, args)
elif args.pgd is True:
imgs, target = pgd_train(imgs, target, model, criterion)
output = model(imgs)
loss = criterion(output, target)
optimizer.zero_grad()
if args.half_prec:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
timer.batch_end()
corr1, corr5 = correct(output.data, target, topk=(1, 5))
reduced_loss, batch_total = loss.data.item(), imgs.size(0)
if args.distributed is True:
metrics = torch.tensor([batch_total, reduced_loss, corr1,
corr5]).float().cuda()
batch_total, reduced_loss, corr1, corr5 = dist_utils.sum_tensor(
metrics).cpu().numpy()
reduced_loss = reduced_loss / dist_utils.env_world_size()
else:
corr1, corr5 = corr1.item(), corr5.item()
top1acc = corr1 * (100.0 / batch_total)
top5acc = corr5 * (100.0 / batch_total)
losses.update(reduced_loss, batch_total)
top1.update(top1acc, batch_total)
top5.update(top5acc, batch_total)
if should_print(batch_num, trn_loader, args) is True:
tb.log_memory()
tb.log_trn_times(timer.batch_time.val, timer.data_time.val,
imgs.size(0))
tb.log_trn_loss(losses.val, top1.val, top5.val)
recv_gbit, transmit_gbit = net_meter.update_bandwidth()
tb.log("sizes/batch_total", batch_total)
tb.log('net/recv_gbit', recv_gbit)
tb.log('net/transmit_gbit', transmit_gbit)
output = (
f'Epoch: [{epoch}][{batch_num}/{len(trn_loader)}]\t'
f'Time {timer.batch_time.val:.3f} ({timer.batch_time.avg:.3f})\t'
f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
f'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
f'Acc@5 {top5.val:.3f} ({top5.avg:.3f})\t'
f'Data {timer.data_time.val:.3f} ({timer.data_time.avg:.3f})\t'
f'BW {recv_gbit:.3f} {transmit_gbit:.3f}')
log.verbose(output)
tb.update_step_count(batch_total)
def do_train(cfg, model, center_criterion, train_loader, val_loader, optimizer,
optimizer_center, scheduler, loss_fn, num_query):
log_period = cfg.SOLVER.LOG_PERIOD
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
eval_period = cfg.SOLVER.EVAL_PERIOD
device = "cuda"
epochs = cfg.SOLVER.MAX_EPOCHS
start_epoch = cfg.SOLVER.START_EPOCH
logger = logging.getLogger("reid_baseline.train")
logger.info('start training')
if device:
if torch.cuda.device_count() > 1:
print('Using {} GPUs for training'.format(
torch.cuda.device_count()))
model = nn.DataParallel(model)
model.to(device)
loss_meter = AverageMeter()
acc_meter1 = AverageMeter()
acc_meter2 = AverageMeter()
# acc_cam = AverageMeter()
evaluator = R1_mAP_eval(num_query,
max_rank=50,
feat_norm=cfg.TEST.FEAT_NORM)
if torch.cuda.device_count() > 1:
model.module.base._freeze_stages()
else:
model.base._freeze_stages()
logger.info('Freezing the stages number:{}'.format(cfg.MODEL.FROZEN))
# train
for epoch in range(start_epoch, epochs + 1):
start_time = time.time()
loss_meter.reset()
acc_meter1.reset()
acc_meter2.reset()
# acc_cam.reset()
evaluator.reset()
scheduler.step()
model.train()
for n_iter, (img, vid, _) in enumerate(train_loader):
optimizer.zero_grad()
optimizer_center.zero_grad()
img = img.to(device)
target = vid.to(device)
# camid = camid.to(device)
scores, feat = model(img, target)
loss = loss_fn(scores, feat, target)
loss.backward()
optimizer.step()
if 'center' in cfg.MODEL.METRIC_LOSS_TYPE:
for param in center_criterion.parameters():
param.grad.data *= (1. / cfg.SOLVER.CENTER_LOSS_WEIGHT)
optimizer_center.step()
acc = [(score.max(1)[1] == target).float().mean()
for score in scores]
# cam_acc = (cam_score.max(1)[1] == camid).float().mean()
loss_meter.update(loss.item(), img.shape[0])
acc_meter1.update(acc[0].item(), 1)
acc_meter2.update(acc[1].item(), 1)
# acc_cam.update(cam_acc.item(), 1)
if (n_iter + 1) % log_period == 0:
logger.info(
"Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Acc1: {:.3f}, Acc2: {:.3f}, Base Lr: {:.2e}"
.format(epoch, (n_iter + 1), len(train_loader),
loss_meter.avg, acc_meter1.avg, acc_meter2.avg,
scheduler.get_lr()[0]))
end_time = time.time()
time_per_batch = (end_time - start_time) / (n_iter + 1)
logger.info(
"Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]"
.format(epoch, time_per_batch,
train_loader.batch_size / time_per_batch))
if epoch % checkpoint_period == 0:
if torch.cuda.device_count() > 1:
torch.save(
{
'static_dict': model.module.state_dict(),
'optimizer_static_dict': optimizer.state_dict()
},
os.path.join(cfg.OUTPUT_DIR,
cfg.MODEL.NAME + '_{}.pth'.format(epoch)))
else:
torch.save(
{
'static_dict': model.state_dict(),
'optimizer_static_dict': optimizer.state_dict()
},
os.path.join(cfg.OUTPUT_DIR,
cfg.MODEL.NAME + '_{}.pth'.format(epoch)))
if epoch % eval_period == 0:
model.eval()
for n_iter, (img, vid, camid, _, _) in enumerate(val_loader):
with torch.no_grad():
img = img.to(device)
feat = model(img)
evaluator.update((feat, vid, camid))
cmc, mAP, _, _, _, _, _ = evaluator.compute()
logger.info("Validation Results - Epoch: {}".format(epoch))
logger.info("mAP: {:.1%}".format(mAP))
for r in [1, 5, 10]:
logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(
r, cmc[r - 1]))
def main(cfg, device):
init_seeds()
cfg.use_fp16 = False if device.type == 'cpu' else cfg.use_fp16
# logging ----------------------------------------------------------------------------------------------------------------------------------------
logger_root = f'Results/{cfg.dataset}'
if not os.path.isdir(logger_root):
os.makedirs(logger_root, exist_ok=True)
logtime = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
result_dir = os.path.join(logger_root, f'{logtime}-{cfg.log}')
# result_dir = os.path.join(logger_root, f'ablation_study-{cfg.log}') #TODO
logger = Logger(logging_dir=result_dir, DEBUG=False)
logger.set_logfile(logfile_name='log.txt')
save_params(cfg, f'{result_dir}/params.json', json_format=True)
logger.debug(f'Result Path: {result_dir}')
# model, optimizer, scheduler --------------------------------------------------------------------------------------------------------------------
opt_lvl = 'O1' if cfg.use_fp16 else 'O0'
n_classes = cfg.n_classes
net1 = ResNet(arch=cfg.net1, num_classes=n_classes, pretrained=True)
optimizer1 = build_sgd_optimizer(net1.parameters(), cfg.lr, cfg.weight_decay)
net1, optimizer1 = amp.initialize(net1.to(device), optimizer1, opt_level=opt_lvl, keep_batchnorm_fp32=None, loss_scale=None, verbosity=0)
net2 = ResNet(arch=cfg.net2, num_classes=n_classes, pretrained=True)
optimizer2 = build_sgd_optimizer(net2.parameters(), cfg.lr, cfg.weight_decay)
net2, optimizer2 = amp.initialize(net2.to(device), optimizer2, opt_level=opt_lvl, keep_batchnorm_fp32=None, loss_scale=None, verbosity=0)
lr_plan = make_lr_plan(cfg.lr, cfg.stage1, cfg.epochs)
with open(f'{result_dir}/network.txt', 'w') as f:
f.writelines(net1.__repr__())
f.write('\n\n---------------------------\n\n')
f.writelines(net1.__repr__())
# drop rate scheduler ----------------------------------------------------------------------------------------------------------------------------
T_k = cfg.stage1
final_drop_rate = 0.25
final_ldl_rate = cfg.ldl_rate
drop_rate_scheduler = np.ones(cfg.epochs) * final_drop_rate
drop_rate_scheduler[:T_k] = np.linspace(0, final_drop_rate, T_k)
drop_rate_scheduler[T_k:cfg.epochs] = np.linspace(final_drop_rate, final_ldl_rate, cfg.epochs - T_k)
# dataset, dataloader ----------------------------------------------------------------------------------------------------------------------------
transform = build_transform(rescale_size=cfg.rescale_size, crop_size=cfg.crop_size)
dataset = build_webfg_dataset(os.path.join(cfg.database, cfg.dataset), CLDataTransform(transform['train']), transform['test'])
logger.debug(f"Number of Training Samples: {dataset['n_train_samples']}")
logger.debug(f"Number of Testing Samples: {dataset['n_test_samples']}")
train_loader = DataLoader(dataset['train'], batch_size=cfg.batch_size, shuffle=True, num_workers=8, pin_memory=True)
test_loader = DataLoader(dataset['test'], batch_size=16, shuffle=False, num_workers=8, pin_memory=True)
# meters -----------------------------------------------------------------------------------------------------------------------------------------
train_loss1, train_loss2 = AverageMeter(), AverageMeter()
train_accuracy1, train_accuracy2 = AverageMeter(), AverageMeter()
iter_time = AverageMeter()
# training ---------------------------------------------------------------------------------------------------------------------------------------
start_epoch = 0
best_accuracy1, best_accuracy2 = 0.0, 0.0
best_epoch1, best_epoch2 = None, None
if cfg.dataset == 'cifar100' and cfg.noise_type != 'clean':
t = torch.tensor(dataset['train'].noisy_labels)
else:
t = torch.tensor(dataset['train'].targets)
labels2learn1 = torch.full(size=(dataset['n_train_samples'], n_classes), fill_value=0.0)
labels2learn1.scatter_(dim=1, index=torch.unsqueeze(t, dim=1), value=1.0 * 10)
labels2learn2 = labels2learn1
flag = [0, 0, 0]
for epoch in range(start_epoch, cfg.epochs):
start_time = time.time()
train_loss1.reset()
train_accuracy1.reset()
train_loss2.reset()
train_accuracy2.reset()
net1.train()
net2.train()
adjust_lr(optimizer1, lr_plan[epoch])
adjust_lr(optimizer2, lr_plan[epoch])
optimizer1.zero_grad()
optimizer2.zero_grad()
# train this epoch
for it, sample in enumerate(train_loader):
s = time.time()
# optimizer1.zero_grad()
# optimizer2.zero_grad()
indices = sample['index']
x1, x2 = sample['data']
x1, x2 = x1.to(device), x2.to(device)
y0 = sample['label'].to(device)
y = get_smoothed_label_distribution(y0, nc=n_classes, epsilon=cfg.epsilon)
output1 = net1(x1)
output2 = net2(x2)
logits1 = output1['logits']
logits2 = output2['logits']
if epoch < cfg.stage1: # warmup
if flag[0] == 0:
step_flagging('stage 1')
flag[0] += 1
loss1 = cross_entropy(logits1, y)
loss2 = cross_entropy(logits2, y)
else: # learn label distributions
if flag[1] == 0:
step_flagging('stage 2')
flag[1] += 1
with torch.no_grad():
cce_losses1 = cross_entropy(logits1, y, reduction='none')
cce_losses2 = cross_entropy(logits2, y, reduction='none')
losses1 = cce_losses1
losses2 = cce_losses2
# ent_losses1 = entropy_loss(logits1, reduction='none')
# ent_losses2 = entropy_loss(logits2, reduction='none')
# losses1 = cce_losses1 + ent_losses1 # (N)
# losses2 = cce_losses2 + ent_losses2 # (N)
sample_selection = sample_selector(losses1, losses2, drop_rate_scheduler[epoch])
# for selected "clean" samples, train in a co-teaching manner
logits_clean1 = logits1[sample_selection['clean2']]
logits_clean2 = logits2[sample_selection['clean1']]
y_clean1 = y[sample_selection['clean2']]
y_clean2 = y[sample_selection['clean1']]
losses_clean1 = cross_entropy(logits_clean1, y_clean1, reduction='none') + entropy_loss(logits_clean1, reduction='none') # (Nc1)
losses_clean2 = cross_entropy(logits_clean2, y_clean2, reduction='none') + entropy_loss(logits_clean2, reduction='none') # (Nc2)
loss_c1_1 = losses_clean1.mean()
loss_c2_1 = losses_clean2.mean()
# for selected "unclean" samples, train in a label distribution learning manner (exchange again)
y_t1 = labels2learn1[indices, :].clone().to(device)
y_t2 = labels2learn2[indices, :].clone().to(device)
y_t1.requires_grad = True
y_t2.requires_grad = True
y_d1 = F.softmax(y_t1, dim=1) + 1e-8
y_d2 = F.softmax(y_t2, dim=1) + 1e-8
logits_unclean1 = logits1[sample_selection['unclean2']]
logits_unclean2 = logits2[sample_selection['unclean1']]
y_d_unclean1 = y_d1[sample_selection['unclean2']]
y_d_unclean2 = y_d2[sample_selection['unclean1']]
w1 = np.random.beta(cfg.phi, cfg.phi, logits_unclean1.size(0))
w2 = np.random.beta(cfg.phi, cfg.phi, logits_unclean2.size(0))
w1 = x1.new(w1).view(logits_unclean1.size(0), 1, 1, 1)
w2 = x2.new(w2).view(logits_unclean2.size(0), 1, 1, 1)
idx1 = np.random.choice(sample_selection['clean2'].cpu().numpy(), logits_unclean1.size(0), replace=False if sample_selection['clean2'].size(0) >= logits_unclean1.size(0) else True)
idx1 = torch.tensor(idx1).to(device)
idx2 = np.random.choice(sample_selection['clean1'].cpu().numpy(), logits_unclean2.size(0), replace=False if sample_selection['clean1'].size(0) >= logits_unclean2.size(0) else True)
idx2 = torch.tensor(idx2).to(device)
mixed_x1 = w1 * x1[sample_selection['unclean2']] + (1-w1) * x1[idx1]
mixed_x2 = w2 * x2[sample_selection['unclean1']] + (1-w2) * x2[idx2]
mixed_y1 = w1 * y_d_unclean1 + (1-w1) * y_d1[idx1]
mixed_y2 = w2 * y_d_unclean2 + (1-w2) * y_d2[idx2]
mixed_output1 = net1(mixed_x1)
mixed_output2 = net2(mixed_x2)
mixed_logits1 = mixed_output1['logits']
mixed_logits2 = mixed_output2['logits']
loss_c1_2 = kl_div(F.softmax(mixed_logits1, dim=1) + 1e-8, mixed_y1).mean()
loss_c2_2 = kl_div(F.softmax(mixed_logits2, dim=1) + 1e-8, mixed_y2).mean()
loss_c1 = loss_c1_1 + loss_c1_2 * cfg.beta
loss_c2 = loss_c2_1 + loss_c2_2 * cfg.beta
# consistency loss
loss_o1 = cross_entropy(F.softmax(y_t1[sample_selection['clean2']], dim=1), y[sample_selection['clean2']])
loss_o2 = cross_entropy(F.softmax(y_t2[sample_selection['clean1']], dim=1), y[sample_selection['clean1']])
# final loss
loss1 = (1 - cfg.alpha) * loss_c1 + cfg.alpha * loss_o1
loss2 = (1 - cfg.alpha) * loss_c2 + cfg.alpha * loss_o2
train_acc1 = accuracy(logits1, y0, topk=(1,))
train_acc2 = accuracy(logits2, y0, topk=(1,))
train_loss1.update(loss1.item(), x1.size(0))
train_loss2.update(loss2.item(), x2.size(0))
train_accuracy1.update(train_acc1[0], x1.size(0))
train_accuracy2.update(train_acc2[0], x2.size(0))
if cfg.use_fp16:
with amp.scale_loss(loss1, optimizer1) as scaled_loss1:
scaled_loss1.backward()
with amp.scale_loss(loss2, optimizer2) as scaled_loss2:
scaled_loss2.backward()
else:
loss1.backward()
loss2.backward()
optimizer1.step()
optimizer2.step()
optimizer1.zero_grad()
optimizer2.zero_grad()
if epoch >= cfg.stage1:
y_t1.data.sub_(cfg.lmd * y_t1.grad.data)
y_t2.data.sub_(cfg.lmd * y_t2.grad.data)
labels2learn1[indices, :] = y_t1.detach().clone().cpu().data
labels2learn2[indices, :] = y_t2.detach().clone().cpu().data
del y_t1, y_t2
iter_time.update(time.time() - s, 1)
if (cfg.log_freq is not None and (it + 1) % cfg.log_freq == 0) or (it + 1 == len(train_loader)):
total_mem = torch.cuda.get_device_properties(0).total_memory / 2**30
mem = torch.cuda.memory_reserved() / 2**30
console_content = f"Epoch:[{epoch + 1:>3d}/{cfg.epochs:>3d}] " \
f"Iter:[{it + 1:>4d}/{len(train_loader):>4d}] " \
f"Train Accuracy 1:[{train_accuracy1.avg:6.2f}] " \
f"Train Accuracy 2:[{train_accuracy2.avg:6.2f}] " \
f"Loss 1:[{train_loss1.avg:4.4f}] " \
f"Loss 2:[{train_loss2.avg:4.4f}] " \
f"GPU-MEM:[{mem:6.3f}/{total_mem:6.3f} Gb] " \
f"{iter_time.avg:6.2f} sec/iter"
logger.debug(console_content)
# evaluate this epoch
test_accuracy1 = evaluate(test_loader, net1, device)
test_accuracy2 = evaluate(test_loader, net2, device)
if test_accuracy1 > best_accuracy1:
best_accuracy1 = test_accuracy1
best_epoch1 = epoch + 1
torch.save(net1.state_dict(), f'{result_dir}/net1_best_epoch.pth')
if test_accuracy2 > best_accuracy2:
best_accuracy2 = test_accuracy2
best_epoch2 = epoch + 1
torch.save(net2.state_dict(), f'{result_dir}/net2_best_epoch.pth')
# logging this epoch
runtime = time.time() - start_time
logger.info(f'epoch: {epoch + 1:>3d} | '
f'train loss(1/2): ({train_loss1.avg:>6.4f}/{train_loss2.avg:>6.4f}) | '
f'train accuracy(1/2): ({train_accuracy1.avg:>6.3f}/{train_accuracy2.avg:>6.3f}) | '
f'test accuracy(1/2): ({test_accuracy1:>6.3f}/{test_accuracy2:>6.3f}) | '
f'epoch runtime: {runtime:6.2f} sec | '
f'best accuracy(1/2): ({best_accuracy1:6.3f}/{best_accuracy2:6.3f}) @ epoch: ({best_epoch1:03d}/{best_epoch2:03d})')
plot_results_cotraining(result_file=f'{result_dir}/log.txt')
torch.save(labels2learn1, f'{result_dir}/labels_learned.pt')
# rename results dir -----------------------------------------------------------------------------------------------------------------------------
best_accuracy = max(best_accuracy1, best_accuracy2)
os.rename(result_dir, f'{result_dir}-bestAcc_{best_accuracy:.4f}')
def do_train(cfg, model, center_criterion, train_loader, val_loader, optimizer,
optimizer_center, scheduler, loss_fn, num_query):
log_period = cfg.SOLVER.LOG_PERIOD
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
eval_period = cfg.SOLVER.EVAL_PERIOD
device = "cuda"
epochs = cfg.SOLVER.MAX_EPOCHS
logger = logging.getLogger("reid_baseline.train")
logger.info('start training')
print("torch.cuda.device_count()", torch.cuda.device_count())
if device:
model.to(device)
if torch.cuda.device_count() > 1:
print('Using {} GPUs for training'.format(
torch.cuda.device_count()))
print("多卡训练")
# model = DDP(model, delay_allreduce=True) # 必须在initialze之后
# model = nn.DataParallel(model)
# model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 字母小写o,不是零。
torch.distributed.init_process_group(
'gloo',
init_method='file:///tmp/somefile',
rank=0,
world_size=1)
# model = convert_syncbn_model(model)
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
# model = DistributedDataParallel(model, delay_allreduce=True)
# model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters=True)
model = nn.DataParallel(model)
# model = convert_syncbn_model(model)
else:
print("单卡训练")
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
model.to(device=0)
loss_meter = AverageMeter()
acc_meter = AverageMeter()
# evaluator = R1_mAP_eval(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)
# model.base._freeze_stages()
logger.info('Freezing the stages number:{}'.format(cfg.MODEL.FROZEN))
# model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
for epoch in range(1, epochs + 1):
if epoch == 5:
print("balance 数据训练")
# cfg.DATASETS.ROOT_DIR = '/home/lab3/bi/0716/Veri/ai_city/tools/mix_train_balance_flip.pkl'
cfg.DATASETS.ROOT_DIR = 'datasets/mix_train_balance.pkl'
train_loader, val_loader, num_query, num_classes = make_dataloader(
cfg)
# model.base._freeze_stages()
start_time = time.time()
loss_meter.reset()
acc_meter.reset()
# evaluator.reset()
scheduler.step()
model.train()
# print(scheduler.get_lr()[0])
for n_iter, (img, vid) in enumerate(tqdm(train_loader)):
optimizer.zero_grad()
optimizer_center.zero_grad()
img = img.to(device)
target = vid.to(device)
#grid mask
# img = grid(img)
# score, feat,score_f1,score_f2,score_f3,f4,f4_score = model(img, target)
# score, feat,score_f1,score_f2,feat1,score_layer2 = model(img, target)
score, feat, score_f1, score_f2, feat1 = model(img, target)
# print(feat.shape)
loss = loss_fn(score, feat, target, score_f1, score_f2, feat1)
# loss = loss_fn(score, feat, target,score_f1,score_f2,feat1,score_layer2)
if cfg.SOLVER.FP16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
# scaled_loss.backward(retain_graph=True)
else:
loss.backward()
# loss.backward()
optimizer.step()
if 'center' in cfg.MODEL.METRIC_LOSS_TYPE:
for param in center_criterion.parameters():
param.grad.data *= (1. / cfg.SOLVER.CENTER_LOSS_WEIGHT)
optimizer_center.step()
acc = (score.max(1)[1] == target).float().mean()
loss_meter.update(loss.item(), img.shape[0])
acc_meter.update(acc, 1)
# print(loss_meter.val)
if (n_iter + 1) % log_period == 0:
logger.info(
"Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Acc: {:.3f}, Base Lr: {:.2e}"
.format(epoch, (n_iter + 1), len(train_loader),
loss_meter.avg, acc_meter.avg,
scheduler.get_lr()[0]))
end_time = time.time()
time_per_batch = (end_time - start_time) / (n_iter + 1)
logger.info(
"Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]"
.format(epoch, time_per_batch,
train_loader.batch_size / time_per_batch))
if epoch % checkpoint_period == 0:
torch.save(
model.state_dict(),
os.path.join(cfg.OUTPUT_DIR,
cfg.MODEL.NAME + '_epoch{}.pth'.format(epoch)))
if epoch == 10:
reduce_model_dict = model.half().state_dict()
del_keys = []
for key in reduce_model_dict.keys():
if 'class' in key or 'sub1' in key or 'sub2' in key or 'base.fc' in key:
del_keys.append(key)
for key in del_keys:
del reduce_model_dict[key]
torch.save(
reduce_model_dict,
os.path.join(
cfg.OUTPUT_DIR, cfg.MODEL.NAME +
str(cfg.INPUT.SIZE_TRAIN[0]) + 'half.pth'))
def do_train(Cfg, model, center_criterion, train_loader, val_loader, optimizer,
optimizer_center, scheduler, loss_fn, num_query):
log_period = Cfg.LOG_PERIOD
checkpoint_period = Cfg.CHECKPOINT_PERIOD
eval_period = Cfg.EVAL_PERIOD
output_dir = Cfg.LOG_DIR
device = "cuda"
epochs = Cfg.MAX_EPOCHS
logger = logging.getLogger('{}.train'.format(Cfg.PROJECT_NAME))
logger.info('start training')
if device:
if torch.cuda.device_count() > 1:
print('Using {} GPUs for training'.format(
torch.cuda.device_count()))
model = nn.DataParallel(model)
model.to(device)
loss_meter = AverageMeter()
acc_meter = AverageMeter()
evaluator = R1_mAP(num_query, max_rank=50, feat_norm=Cfg.FEAT_NORM)
#train
for epoch in range(1, epochs + 1):
start_time = time.time()
loss_meter.reset()
acc_meter.reset()
evaluator.reset()
model.train()
for iter, (img, vid) in enumerate(train_loader):
optimizer.zero_grad()
optimizer_center.zero_grad()
img = img.to(device)
target = vid.to(device)
score, feat = model(img, target)
loss = loss_fn(score, feat, target)
loss.backward()
optimizer.step()
if 'center' in Cfg.LOSS_TYPE:
for param in center_criterion.parameters():
param.grad.data *= (1. / Cfg.CENTER_LOSS_WEIGHT)
optimizer_center.step()
acc = (score.max(1)[1] == target).float().mean()
loss_meter.update(loss.item(), img.shape[0])
acc_meter.update(acc, 1)
if (iter + 1) % log_period == 0:
logger.info(
"Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Acc: {:.3f}, Base Lr: {:.2e}"
.format(epoch, (iter + 1), len(train_loader),
loss_meter.avg, acc_meter.avg,
scheduler.get_lr()[0]))
end_time = time.time()
time_per_batch = (end_time - start_time) / (iter + 1)
logger.info(
"Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]"
.format(epoch, time_per_batch,
train_loader.batch_size / time_per_batch))
scheduler.step()
if epoch % checkpoint_period == 0:
torch.save(model.state_dict(),
output_dir + Cfg.MODEL_NAME + '_{}.pth'.format(epoch))
if epoch % eval_period == 0:
model.eval()
for iter, (img, vid, camid) in enumerate(val_loader):
with torch.no_grad():
img = img.to(device)
feat = model(img)
evaluator.update((feat, vid, camid))
cmc, mAP, _, _, _, _ = evaluator.compute()
logger.info("Validation Results - Epoch: {}".format(epoch))
logger.info("mAP: {:.1%}".format(mAP))
for r in [1, 5, 10]:
logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(
r, cmc[r - 1]))
def do_train(Cfg, model_G, model_Dip, model_Dii, model_D_reid, train_loader,
val_loader, optimizerG, optimizerDip, optimizerDii, GAN_loss,
L1_loss, ReID_loss, schedulerG, schedulerDip, schedulerDii):
log_period = Cfg.SOLVER.LOG_PERIOD
checkpoint_period = Cfg.SOLVER.CHECKPOINT_PERIOD
eval_period = Cfg.SOLVER.EVAL_PERIOD
output_dir = Cfg.DATALOADER.LOG_DIR
# need modified the following in cfg
epsilon = 0.00001
margin = 0.4
####################################
device = "cuda"
epochs = Cfg.SOLVER.MAX_EPOCHS
logger = logging.getLogger('pose-transfer-gan.train')
logger.info('Start training')
if device:
if torch.cuda.device_count() > 1:
print('Using {} GPUs for training'.format(
torch.cuda.device_count()))
model_G = nn.DataParallel(model_G)
model_Dii = nn.DataParallel(model_Dii)
model_Dip = nn.DataParallel(model_Dip)
model_G.to(device)
model_Dip.to(device)
model_Dii.to(device)
model_D_reid.to(device)
lossG_meter = AverageMeter()
lossDip_meter = AverageMeter()
lossDii_meter = AverageMeter()
distDreid_meter = AverageMeter()
fake_ii_pool = ImagePool(50)
fake_ip_pool = ImagePool(50)
#evaluator = R1_mAP(num_query, max_rank=50, feat_norm=Cfg.TEST.FEAT_NORM)
#train
for epoch in range(1, epochs + 1):
start_time = time.time()
lossG_meter.reset()
lossDip_meter.reset()
lossDii_meter.reset()
distDreid_meter.reset()
schedulerG.step()
schedulerDip.step()
schedulerDii.step()
model_G.train()
model_Dip.train()
model_Dii.train()
model_D_reid.eval()
for iter, batch in enumerate(train_loader):
img1 = batch['img1'].to(device)
pose1 = batch['pose1'].to(device)
img2 = batch['img2'].to(device)
pose2 = batch['pose2'].to(device)
input_G = (img1, pose2)
#forward
fake_img2 = model_G(input_G)
optimizerG.zero_grad()
#train G
input_Dip = torch.cat((fake_img2, pose2), 1)
pred_fake_ip = model_Dip(input_Dip)
loss_G_ip = GAN_loss(pred_fake_ip, True)
input_Dii = torch.cat((fake_img2, img1), 1)
pred_fake_ii = model_Dii(input_Dii)
loss_G_ii = GAN_loss(pred_fake_ii, True)
loss_L1, _, _ = L1_loss(fake_img2, img2)
feats_real = model_D_reid(img2)
feats_fake = model_D_reid(fake_img2)
dist_cos = torch.acos(
torch.clamp(torch.sum(feats_real * feats_fake, 1),
-1 + epsilon, 1 - epsilon))
same_id_tensor = torch.FloatTensor(
dist_cos.size()).fill_(1).to('cuda')
dist_cos_margin = torch.max(dist_cos - margin,
torch.zeros_like(dist_cos))
loss_reid = ReID_loss(dist_cos_margin, same_id_tensor)
factor = loss_reid_factor(epoch)
loss_G = 0.5 * loss_G_ii * Cfg.LOSS.GAN_WEIGHT + 0.5 * loss_G_ip * Cfg.LOSS.GAN_WEIGHT + loss_L1 + loss_reid * Cfg.LOSS.REID_WEIGHT * factor
loss_G.backward()
optimizerG.step()
#train Dip
for i in range(Cfg.SOLVER.DG_RATIO):
optimizerDip.zero_grad()
real_input_ip = torch.cat((img2, pose2), 1)
fake_input_ip = fake_ip_pool.query(
torch.cat((fake_img2, pose2), 1).data)
pred_real_ip = model_Dip(real_input_ip)
loss_Dip_real = GAN_loss(pred_real_ip, True)
pred_fake_ip = model_Dip(fake_input_ip)
loss_Dip_fake = GAN_loss(pred_fake_ip, False)
loss_Dip = 0.5 * Cfg.LOSS.GAN_WEIGHT * (loss_Dip_real +
loss_Dip_fake)
loss_Dip.backward()
optimizerDip.step()
#train Dii
for i in range(Cfg.SOLVER.DG_RATIO):
optimizerDii.zero_grad()
real_input_ii = torch.cat((img2, img1), 1)
fake_input_ii = fake_ii_pool.query(
torch.cat((fake_img2, img1), 1).data)
pred_real_ii = model_Dii(real_input_ii)
loss_Dii_real = GAN_loss(pred_real_ii, True)
pred_fake_ii = model_Dii(fake_input_ii)
loss_Dii_fake = GAN_loss(pred_fake_ii, False)
loss_Dii = 0.5 * Cfg.LOSS.GAN_WEIGHT * (loss_Dii_real +
loss_Dii_fake)
loss_Dii.backward()
optimizerDii.step()
lossG_meter.update(loss_G.item(), 1)
lossDip_meter.update(loss_Dip.item(), 1)
lossDii_meter.update(loss_Dii.item(), 1)
distDreid_meter.update(dist_cos.mean().item(), 1)
if (iter + 1) % log_period == 0:
logger.info(
"Epoch[{}] Iteration[{}/{}] G Loss: {:.3f}, Dip Loss: {:.3f}, Dii Loss: {:.3f}, Base G_Lr: {:.2e}, Base Dip_Lr: {:.2e}, Base Dii_Lr: {:.2e}"
.format(epoch, (iter + 1), len(train_loader),
lossG_meter.avg, lossDip_meter.avg,
lossDii_meter.avg,
schedulerG.get_lr()[0],
schedulerDip.get_lr()[0],
schedulerDii.get_lr()[0])) #scheduler.get_lr()[0]
logger.info("ReID Cos Distance: {:.3f}".format(
distDreid_meter.avg))
end_time = time.time()
time_per_batch = (end_time - start_time) / (iter + 1)
logger.info(
"Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]"
.format(epoch, time_per_batch,
train_loader.batch_size / time_per_batch))
if epoch % checkpoint_period == 0:
torch.save(model_G.state_dict(),
output_dir + 'model_G_{}.pth'.format(epoch))
torch.save(model_Dip.state_dict(),
output_dir + 'model_Dip_{}.pth'.format(epoch))
torch.save(model_Dii.state_dict(),
output_dir + 'model_Dii_{}.pth'.format(epoch))
#
if epoch % eval_period == 0:
np.save(output_dir + 'train_Bx6x128x64_epoch{}.npy'.format(epoch),
fake_ii_pool.images[0].cpu().numpy())
logger.info('Entering Evaluation...')
tmp_results = []
model_G.eval()
for iter, batch in enumerate(val_loader):
with torch.no_grad():
img1 = batch['img1'].to(device)
pose1 = batch['pose1'].to(device)
img2 = batch['img2'].to(device)
pose2 = batch['pose2'].to(device)
input_G = (img1, pose2)
fake_img2 = model_G(input_G)
tmp_result = torch.cat((img1, img2, fake_img2),
1).cpu().numpy()
tmp_results.append(tmp_result)
np.save(output_dir + 'test_Bx6x128x64_epoch{}.npy'.format(epoch),
tmp_results[0])
def train(train_loader, model, criterion, optimizer, epoch, rank, args, logger, cfg):
batch_time = AverageMeter('Time', ':6.3f')
data_time = AverageMeter('Data', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
top5 = AverageMeter('Acc@5', ':6.2f')
progress = ProgressMeter(len(train_loader), [batch_time, data_time, losses, top1, top5],
prefix="Epoch: [{}]".format(epoch))
# switch to train mode
model.train()
end = time.time()
for iter, traindata in enumerate(train_loader):
data, label, _, meta = traindata
if isinstance(data, (list,)):
for i in range(len(data)):
data[i] = data[i].cuda(rank)
else:
data = data.cuda(rank)
label = label.cuda(rank)
for key, val in meta.items():
if isinstance(val, (list,)):
for i in range(len(val)):
val[i] = val[i].cuda(non_blocking=True)
else:
meta[key] = val.cuda(non_blocking=True)
# measure data loading time
data_time.update(time.time() - end)
# compute output
if cfg.DETECTION.ENABLE:
output = model(data, meta["boxes"])
else:
output = model(data)
loss = criterion(output, label)
adjust_lr(optimizer, epoch, cfg.SOLVER.BASE_LR)
# measure accuracy and record loss
acc1, acc5 = topks_correct(output, label, (1, 5))
torch.distributed.barrier()
reduced_loss = reduce_mean(loss, args.nprocs)
reduced_acc1 = reduce_mean(acc1, args.nprocs)
reduced_acc5 = reduce_mean(acc5, args.nprocs)
losses.update(reduced_loss.item(), data[0].size(0))
top1.update(reduced_acc1.item(), data[0].size(0))
top5.update(reduced_acc5.item(), data[0].size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
# with amp.scale_loss(loss, optimizer) as scaled_loss:
# scaled_loss.backward()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if iter % args.print_freq == 0:
train_message = progress.display(iter)
logger.info('Train Phase:' + train_message)
return losses.avg
def validate(self, val_loader, tb_logger=None):
batch_time = AverageMeter(0)
loss_meter = AverageMeter(0)
top1_meter = AverageMeter(0)
top5_meter = AverageMeter(0)
self.model.eval()
criterion = nn.CrossEntropyLoss()
end = time.time()
with torch.no_grad():
for batch_idx, (x, y) in enumerate(val_loader):
x, y = x.cuda(), y.cuda()
num = x.size(0)
out = self.model(x)
loss = criterion(out, y)
top1, top5 = accuracy(out, y, top_k=(1, 5))
loss_meter.update(loss.item(), num)
top1_meter.update(top1.item(), num)
top5_meter.update(top5.item(), num)
batch_time.update(time.time() - end)
end = time.time()
if batch_idx % self.config.logging.print_freq == 0:
self._info('Test: [{0}/{1}]\tTime {batch_time.val:.3f} ({batch_time.avg:.3f})'
.format(batch_idx, len(val_loader), batch_time=batch_time))
total_num = torch.tensor([loss_meter.count]).cuda()
loss_sum = torch.tensor([loss_meter.avg * loss_meter.count]).cuda()
top1_sum = torch.tensor([top1_meter.avg * top1_meter.count]).cuda()
top5_sum = torch.tensor([top5_meter.avg * top5_meter.count]).cuda()
dist.all_reduce(total_num)
dist.all_reduce(loss_sum)
dist.all_reduce(top1_sum)
dist.all_reduce(top5_sum)
val_loss = loss_sum.item() / total_num.item()
val_top1 = top1_sum.item() / total_num.item()
val_top5 = top5_sum.item() / total_num.item()
self._info('Prec@1 {:.3f}\tPrec@5 {:.3f}\tLoss {:.3f}\ttotal_num={}'
.format(val_top1, val_top5, val_loss, loss_meter.count))
if dist.is_master():
if val_top1 > self.best_top1:
self.best_top1 = val_top1
if tb_logger is not None:
tb_logger.add_scalar('loss_val', val_loss, self.cur_step)
tb_logger.add_scalar('acc1_val', val_top1, self.cur_step)
tb_logger.add_scalar('acc5_val', val_top5, self.cur_step)
def validate(val_loader, model, criterion, rank, args, logger, cfg):
batch_time = AverageMeter('Time', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
top5 = AverageMeter('Acc@5', ':6.2f')
progress = ProgressMeter(len(val_loader), [batch_time, losses, top1, top5], prefix='Test: ')
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
for iter, valdata in enumerate(val_loader):
data, label, _, meta= valdata
if isinstance(data, (list,)):
for i in range(len(data)):
data[i] = data[i].cuda(rank)
else:
data = data.cuda(rank)
label = label.cuda(rank)
for key, val in meta.items():
if isinstance(val, (list,)):
for i in range(len(val)):
val[i] = val[i].cuda(rank)
else:
meta[key] = val.cuda(rank)
# compute output
if cfg.DETECTION.ENABLE:
output = model(data, meta["boxes"])
else:
output = model(data)
loss = criterion(output, label)
# measure accuracy and record loss
acc1, acc5 = topks_correct(output, label, (1, 5))
# torch.distributed.barrier()
reduced_loss = reduce_mean(loss, args.nprocs)
reduced_acc1 = reduce_mean(acc1, args.nprocs)
reduced_acc5 = reduce_mean(acc5, args.nprocs)
losses.update(reduced_loss.item(), data[0].size(0))
top1.update(reduced_acc1.item(), data[0].size(0))
top5.update(reduced_acc5.item(), data[0].size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if iter % args.print_freq == 0:
val_message = progress.display(iter)
logger.info('Val Phase:' + val_message)
# TODO: this should also be done with the ProgressMeter
logger.info(' * Val Acc@1 {top1.avg:.3f} Val Acc@5 {top5.avg:.3f}'.format(top1=top1, top5=top5))
return top1.avg
def do_train(cfg, model, train_loader, optimizer, scheduler, loss_fn):
log_period = cfg.SOLVER.LOG_PERIOD
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
device = "cuda"
epochs = cfg.SOLVER.MAX_EPOCHS
logger = logging.getLogger("reid_baseline.train")
logger.info('start training')
if device:
model.to(device)
if torch.cuda.device_count() > 1:
print('Using {} GPUs for training'.format(
torch.cuda.device_count()))
model = nn.DataParallel(model)
loss_meter = AverageMeter()
acc_meter = AverageMeter()
# train
scaler = GradScaler()
for epoch in range(1, epochs + 1):
start_time = time.time()
loss_meter.reset()
acc_meter.reset()
model.train()
for n_iter, (img, vid) in enumerate(train_loader):
optimizer.zero_grad()
if cfg.INPUT.AUGMIX:
bs = img[0].size(0)
images_cat = torch.cat(img, dim=0).to(
device) # [3 * batch, 3, 32, 32]
target = vid.to(device)
with autocast():
logits, feat = model(images_cat, target)
logits_orig, logits_augmix1, logits_augmix2 = logits[:bs], logits[
bs:2 * bs], logits[2 * bs:]
loss = loss_fn(logits_orig, feat, target)
p_orig, p_augmix1, p_augmix2 = F.softmax(
logits_orig,
dim=-1), F.softmax(logits_augmix1,
dim=-1), F.softmax(logits_augmix2,
dim=-1)
# Clamp mixture distribution to avoid exploding KL divergence
p_mixture = torch.clamp(
(p_orig + p_augmix1 + p_augmix2) / 3., 1e-7, 1).log()
loss += 12 * (
F.kl_div(p_mixture, p_orig, reduction='batchmean') +
F.kl_div(p_mixture, p_augmix1, reduction='batchmean') +
F.kl_div(p_mixture, p_augmix2,
reduction='batchmean')) / 3.
else:
img = img.to(device)
target = vid.to(device)
with autocast():
if cfg.MODEL.CHANNEL_HEAD:
score, feat, channel_head_feature = model(img, target)
#print(feat.shape, channel_head_feature.shape)
loss = loss_fn(score, feat, channel_head_feature,
target)
else:
score, feat = model(img, target)
loss = loss_fn(score, feat, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
acc = (score.max(1)[1] == target).float().mean()
loss_meter.update(loss.item(), img.shape[0])
acc_meter.update(acc, 1)
if (n_iter + 1) % log_period == 0:
logger.info(
"Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Acc: {:.3f}, Base Lr: {:.2e}"
.format(epoch, (n_iter + 1), len(train_loader),
loss_meter.avg, acc_meter.avg,
scheduler.get_lr()[0]))
scheduler.step()
end_time = time.time()
time_per_batch = (end_time - start_time) / (n_iter + 1)
logger.info(
"Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]"
.format(epoch, time_per_batch,
train_loader.batch_size / time_per_batch))
if epoch % checkpoint_period == 0:
torch.save(
model.state_dict(),
os.path.join(cfg.OUTPUT_DIR,
cfg.MODEL.NAME + '_{}.pth'.format(epoch)))
def evaluate(self, method, epoch):
self.model.eval()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
with torch.no_grad():
if self.train_precision == 'fp16':
self.model.half()
for i, data in enumerate(tqdm(self.test_loader)):
if conf.get()['data']['dali']['avail']:
timages = data[0]["data"]
tlabels = data[0]["label"].squeeze().long()
else:
timages = data[0].type(torch.FloatTensor)
tlabels = data[1].type(torch.LongTensor)
if self.conf.get()['cuda']['avail']:
timages, tlabels = timages.to(self.device), tlabels.to(
self.device)
if self.train_precision == 'fp16':
timages = timages.half()
toutputs = self.model(timages)
tloss = self.criterion(toutputs, tlabels)
ttop1, ttop5 = self.accuracy(toutputs, tlabels, (1, 5))
losses.update(tloss.item(), timages.size(0))
top1.update(ttop1.item(), timages.size(0))
top5.update(ttop5.item(), timages.size(0))
if self.best_acc_top1 < top1.avg:
self.best_acc_top1 = top1.avg
if self.best_acc_top5 < top5.avg:
self.best_acc_top5 = top5.avg
self.save_checkpoint(method, "best")
if self.best_test_loss > losses.avg:
self.best_test_loss = losses.avg
print(
'[{:d}/{:d}] <<<TEST>>> loss({:.4f}) top1({:.4f}) top5({:.4f}) best-top1({:.4f}) best-top5({:.4f})'
.format(epoch + 1, self.train_epochs, losses.avg, top1.avg,
top5.avg, self.best_acc_top1, self.best_acc_top5))
def validate(val_loader, model, criterion, epoch, start_time):
"""Run the validation set through a trained classifier.
Params:
val_loader: Loader for validation set.
model: Classifier instance.
criterion: Loss function.
epoch: Current training epoch.
start_time: Time when training started.
"""
timer = TimeMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
model.eval()
eval_start_time = time.time()
for i, (imgs, target) in enumerate(val_loader):
imgs, target = imgs.cuda(), target.cuda()
# terminates epoch early
if args.short_epoch and (i > 10): break
batch_num = i + 1
timer.batch_start()
if args.distributed:
top1acc, top5acc, loss, batch_total = distributed_predict(
imgs, target, model, criterion)
else:
with torch.no_grad():
output = model(imgs)
loss = criterion(output, target).data
batch_total = imgs.size(0)
top1acc, top5acc = accuracy(output.data, target, topk=(1, 5))
# Eval batch done. Logging results
timer.batch_end()
losses.update(loss, batch_total)
top1.update(top1acc, batch_total)
top5.update(top5acc, batch_total)
if should_print(batch_num, val_loader, args) is True:
output = (
f'Test: [{epoch}][{batch_num}/{len(val_loader)}]\t'
f'Time {timer.batch_time.val:.3f} ({timer.batch_time.avg:.3f})\t'
f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
f'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
f'Acc@5 {top5.val:.3f} ({top5.avg:.3f})')
log.verbose(output)
tb.log_eval(top1.avg, top5.avg, time.time() - eval_start_time)
tb.log('epoch', epoch)
return top1.avg, top5.avg
def train_once(runner):
for epoch in range(runner.train_epochs):
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
runner.model.train()
for i, data in enumerate(runner.tr_loader):
images, labels = data
if conf.get()['cuda']['avail']:
images, labels = images.to(runner.device), labels.to(
runner.device)
runner.model = runner.model.to(runner.device)
runner.optimizer.zero_grad()
outputs, loss = runner.regularizer(images, labels)
loss.backward()
runner.optimizer.step()
ttop1, ttop5 = runner.accuracy(outputs, labels, (1, 5))
losses.update(loss.item(), images.size(0))
top1.update(ttop1.item(), images.size(0))
top5.update(ttop5.item(), images.size(0))
print(
'[{:d}/{:d}] <<<TRAIN>>> lr({:.10f}) loss({:.4f}) top1({:.3f}) top5({:.3f})'
.format(epoch + 1, runner.train_epochs,
runner.optimizer.param_groups[0]['lr'], losses.avg,
top1.avg, top5.avg))
runner.scheduler.step()
def do_train(cfg, model, center_criterion, train_loader, val_loader, optimizer,
optimizer_center, scheduler, loss_fn, num_query, writer):
log_period = cfg.SOLVER.LOG_PERIOD
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
device = "cuda"
epochs = cfg.SOLVER.MAX_EPOCHS
tmp_input_data = torch.rand(
(10, 3, cfg.INPUT.SIZE_TRAIN[0], cfg.INPUT.SIZE_TRAIN[1]))
writer.add_graph(model, (tmp_input_data, ))
logger = logging.getLogger("reid_baseline.train")
logger.info('start training')
if device:
model.to(device)
if torch.cuda.device_count() > 1:
print('Using {} GPUs for training'.format(
torch.cuda.device_count()))
model = nn.DataParallel(model)
model = model.cuda()
else:
if cfg.SOLVER.FP16:
model, optimizer = amp.initialize(model,
optimizer,
opt_level='O1')
loss_meter = AverageMeter()
id_loss_meter = AverageMeter()
tri_loss_meter = AverageMeter()
cen_loss_meter = AverageMeter()
acc_meter = AverageMeter()
lr_meter = AverageMeter()
if cfg.SOLVER.SWA:
swa_model = torch.optim.swa_utils.AveragedModel(model)
# train
for epoch in range(1, epochs + 1):
start_time = time.time()
loss_meter.reset()
acc_meter.reset()
id_loss_meter.reset()
tri_loss_meter.reset()
cen_loss_meter.reset()
lr_meter.reset()
model.train()
if cfg.SOLVER.GRADUAL_UNLOCK:
model.base.gradual_unlock(cfg.SOLVER.MAX_EPOCHS, epoch)
for n_iter, (img, vid) in enumerate(train_loader):
optimizer.zero_grad()
optimizer_center.zero_grad()
img = img.to(device)
target = vid.to(device)
if cfg.DATASETS.MIXUP:
img, target_a, target_b, lam = mixup_data(img, target)
score, feat = model(img, target)
if cfg.DATASETS.MIXUP:
all_loss = mixup_criterion(loss_fn, score, feat, target_a,
target_b, lam)
else:
all_loss = loss_fn(score, feat, target)
loss, id_loss, tri_loss, cen_loss = all_loss
if cfg.SOLVER.FP16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
if 'center' in cfg.MODEL.METRIC_LOSS_TYPE:
for param in center_criterion.parameters():
param.grad.data *= (1. / cfg.SOLVER.CENTER_LOSS_WEIGHT)
optimizer_center.step()
loss_meter.update(loss.item(), img.shape[0])
id_loss_meter.update(id_loss.item(), img.shape[0])
if torch.is_tensor(tri_loss):
tri_loss_meter.update(tri_loss.item(), img.shape[0])
else:
tri_loss_meter.update(tri_loss, 1)
if torch.is_tensor(cen_loss):
cen_loss_meter.update(cen_loss.item(), img.shape[0])
else:
cen_loss_meter.update(cen_loss, 1)
acc = (score.max(1)[1] == target).float().mean()
acc_meter.update(acc, 1)
lr_meter.update(scheduler.get_last_lr()[0])
writer.add_scalar('data/total_loss', loss_meter.avg,
(epoch - 1) * len(train_loader) + n_iter)
writer.add_scalar('data/id_loss', id_loss_meter.avg,
(epoch - 1) * len(train_loader) + n_iter)
writer.add_scalar('data/tri_loss', tri_loss_meter.avg,
(epoch - 1) * len(train_loader) + n_iter)
writer.add_scalar('data/cen_loss', cen_loss_meter.avg,
(epoch - 1) * len(train_loader) + n_iter)
writer.add_scalar('data/learning_rate', lr_meter.avg,
(epoch - 1) * len(train_loader) + n_iter)
if (n_iter + 1) % log_period == 0:
logger.info(
"Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Acc: {:.3f}, Base Lr: {:.2e}"
.format(epoch, (n_iter + 1), len(train_loader),
loss_meter.avg, acc_meter.avg,
scheduler.get_last_lr()[0]))
scheduler.step()
end_time = time.time()
time_per_batch = (end_time - start_time) / (n_iter + 1)
logger.info(
"Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]"
.format(epoch, time_per_batch,
train_loader.batch_size / time_per_batch))
if epoch % checkpoint_period == 0:
src_path = os.path.join(cfg.OUTPUT_DIR,
cfg.MODEL.NAME + '_{}.pth'.format(epoch))
torch.save(model.state_dict(), src_path)
try:
dest_root = os.path.join(
'/mnt/nfs-internstorage/user/zjf/NAIC2020/models',
cfg.SAVE_FLAG)
if not os.path.exists(dest_root):
os.mkdir(dest_root)
dst_path = os.path.join(
dest_root, cfg.MODEL.NAME + '_{}.pth'.format(epoch))
shutil.copy(src_path, dst_path)
except:
print('No bak models...')
pass
if cfg.SOLVER.SWA and epoch in cfg.SOLVER.SWA_START:
swa_model.update_parameters(model)
logger.info('swa combine the {} epoch model'.format(epoch))
if cfg.SOLVER.SWA:
try:
swa_model.cpu()
torch.optim.swa_utils.update_bn(train_loader, swa_model)
swa_model.cuda()
src_path = os.path.join(cfg.OUTPUT_DIR,
cfg.MODEL.NAME + '_swa.pth')
torch.save(swa_model.state_dict(), src_path)
logger.info('swa model is successfuly saved.')
except:
logger.info('swa model save failed.')
def train(runner):
num_iter = 1
if conf.get()['model']['training']['type'] == 'iterative_train':
num_iter = 4
runner.train_epochs = runner.train_epochs // num_iter
for iter_train in range(num_iter):
runner.reset_average_meters()
for epoch in range(runner.train_epochs):
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
runner.model.train()
for i, data in enumerate(tqdm(runner.tr_loader)):
if runner.train_precision == 'fp16':
runner.to_half()
if conf.get()['data']['dali']['avail']:
images = data[0]["data"]
labels = data[0]["label"].squeeze().long()
else:
images, labels = data
if runner.train_precision == 'fp16':
images = images.half()
if conf.get()['cuda']['avail']:
images, labels = images.to(runner.device), labels.to(
torch.device(runner.device))
runner.model = runner.model.to(torch.device(runner.device))
runner.optimizer.zero_grad()
outputs, loss = runner.regularizer(images, labels)
loss.backward()
if runner.train_precision == 'fp16':
runner.to_half()
runner.optimizer.step()
ttop1, ttop5 = runner.accuracy(outputs, labels, (1, 5))
losses.update(loss.item(), images.size(0))
top1.update(ttop1.item(), images.size(0))
top5.update(ttop5.item(), images.size(0))
print(
'[{:d}/{:d}] <<<TRAIN>>> lr({:.10f}) loss({:.4f}) top1({:.3f}) top5({:.3f})'
.format(epoch + 1, runner.train_epochs,
runner.optimizer.param_groups[0]['lr'], losses.avg,
top1.avg, top5.avg))
runner.evaluate("train", epoch)
runner.scheduler.step()
if conf.get()['data']['dali']['avail']:
runner.tr_loader.reset()
runner.test_loader.reset()
if conf.get()['model']['training']['type'] == 'iterative_train':
runner.tr_loader, _, runner.test_loader = get_dataloader(conf)
print('train() Done.')
