def validate(model, data_loader):
print('validating ... ', flush=True, end='')
val_loss_meter = pyutils.AverageMeter('loss1', 'loss2')
model.eval()
with torch.no_grad():
for pack in data_loader:
img = pack['img']
label = pack['label']
if use_gpu:
label = label.cuda(non_blocking=True)
x = model(img)
loss1 = F.multilabel_soft_margin_loss(x, label)
val_loss_meter.add({'loss1': loss1.item()})
model.train()
print('loss: %.4f' % (val_loss_meter.pop('loss1')))
return
def run(args):
assert args.voc12_root is not None
assert args.class_label_dict_path is not None
assert args.train_list is not None
assert args.val_list is not None
assert args.cam_weights_name is not None
assert args.cam_network is not None
assert args.cam_num_epoches is not None
assert args.cam_network_module is not None
model = getattr(importlib.import_module(args.cam_network_module),
args.cam_network)(num_classes=args.num_classes)
train_dataset = dataloader.VOC12ClassificationDataset(
args.train_list,
voc12_root=args.voc12_root,
resize_long=(320, 640),
hor_flip=True,
crop_size=512,
crop_method="random",
class_label_dict_path=args.class_label_dict_path)
train_data_loader = DataLoader(train_dataset,
batch_size=args.cam_batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
max_step = (len(train_dataset) //
args.cam_batch_size) * args.cam_num_epoches
val_dataset = dataloader.VOC12ClassificationDataset(
args.val_list,
voc12_root=args.voc12_root,
crop_size=512,
class_label_dict_path=args.class_label_dict_path)
val_data_loader = DataLoader(val_dataset,
batch_size=args.cam_batch_size,
shuffle=False,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
param_groups = model.trainable_parameters()
optimizer = torchutils.PolyOptimizer([
{
'params': param_groups[0],
'lr': args.cam_learning_rate,
'weight_decay': args.cam_weight_decay
},
{
'params': param_groups[1],
'lr': 10 * args.cam_learning_rate,
'weight_decay': args.cam_weight_decay
},
],
lr=args.cam_learning_rate,
weight_decay=args.cam_weight_decay,
max_step=max_step)
if use_gpu:
model = torch.nn.DataParallel(model).cuda()
model.train()
avg_meter = pyutils.AverageMeter()
timer = pyutils.Timer()
for ep in range(args.cam_num_epoches):
print('Epoch %d/%d' % (ep + 1, args.cam_num_epoches))
correct = 0.
total = 0.
for step, pack in tqdm(enumerate(train_data_loader),
total=len(train_dataset) //
args.cam_batch_size):
img = pack['img']
label = pack['label']
if use_gpu:
label = label.cuda(non_blocking=True)
x = model(img)
_, predicted = torch.max(x.data, 1)
_, actual = torch.max(label.data, 1)
correct += (predicted == actual).sum()
total += label.shape[0]
loss = F.multilabel_soft_margin_loss(x, label)
avg_meter.add({'loss1': loss.item()})
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (optimizer.global_step - 1) % 100 == 0:
acc = 100 * correct / total
timer.update_progress(optimizer.global_step / max_step)
print('step:%5d/%5d' % (optimizer.global_step - 1, max_step),
'loss:%.4f' % (avg_meter.pop('loss1')),
'imps:%.1f' % ((step + 1) * args.cam_batch_size /
timer.get_stage_elapsed()),
'lr: %.4f' % (optimizer.param_groups[0]['lr']),
'etc:%s' % (timer.str_estimated_complete()),
'acc:%s' % acc,
flush=True)
else:
validate(model, val_data_loader)
timer.reset_stage()
try:
state_dict = model.module.state_dict()
except:
state_dict = model.state_dict()
torch.save(state_dict, args.cam_weights_name + '.pth')
if use_gpu:
torch.cuda.empty_cache()
def run(args):
assert args.voc12_root is not None
assert args.class_label_dict_path is not None
assert args.train_list is not None
assert args.ir_label_out_dir is not None
assert args.infer_list is not None
assert args.irn_network is not None
assert args.irn_network_module is not None
path_index = indexing.PathIndex(radius=10,
default_size=(args.irn_crop_size // 4,
args.irn_crop_size // 4))
model = getattr(importlib.import_module(args.irn_network_module),
args.irn_network + 'AffinityDisplacementLoss')(path_index)
train_dataset = dataloader.VOC12AffinityDataset(
args.train_list,
label_dir=args.ir_label_out_dir,
voc12_root=args.voc12_root,
indices_from=path_index.src_indices,
indices_to=path_index.dst_indices,
hor_flip=True,
crop_size=args.irn_crop_size,
crop_method="random",
rescale=(0.5, 1.5))
train_data_loader = DataLoader(train_dataset,
batch_size=args.irn_batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
max_step = (len(train_dataset) //
args.irn_batch_size) * args.irn_num_epoches
param_groups = model.trainable_parameters()
optimizer = torchutils.PolyOptimizer([{
'params': param_groups[0],
'lr': 1 * args.irn_learning_rate,
'weight_decay': args.irn_weight_decay
}, {
'params': param_groups[1],
'lr': 10 * args.irn_learning_rate,
'weight_decay': args.irn_weight_decay
}],
lr=args.irn_learning_rate,
weight_decay=args.irn_weight_decay,
max_step=max_step)
if use_gpu:
model = torch.nn.DataParallel(model).cuda()
model.train()
avg_meter = pyutils.AverageMeter()
timer = pyutils.Timer()
for ep in range(args.irn_num_epoches):
print('Epoch %d/%d' % (ep + 1, args.irn_num_epoches))
for iter, pack in tqdm(enumerate(train_data_loader),
total=len(train_dataset) //
args.irn_batch_size):
img = pack['img']
bg_pos_label = pack['aff_bg_pos_label']
fg_pos_label = pack['aff_fg_pos_label']
neg_label = pack['aff_neg_label']
if use_gpu:
img = img.cuda(non_blocking=True)
bg_pos_label = bg_pos_label.cuda(non_blocking=True)
fg_pos_label = fg_pos_label.cuda(non_blocking=True)
neg_label = neg_label.cuda(non_blocking=True)
pos_aff_loss, neg_aff_loss, dp_fg_loss, dp_bg_loss = model(
img, True)
bg_pos_aff_loss = torch.sum(
bg_pos_label * pos_aff_loss) / (torch.sum(bg_pos_label) + 1e-5)
fg_pos_aff_loss = torch.sum(
fg_pos_label * pos_aff_loss) / (torch.sum(fg_pos_label) + 1e-5)
pos_aff_loss = bg_pos_aff_loss / 2 + fg_pos_aff_loss / 2
neg_aff_loss = torch.sum(
neg_label * neg_aff_loss) / (torch.sum(neg_label) + 1e-5)
dp_fg_loss = torch.sum(dp_fg_loss * torch.unsqueeze(
fg_pos_label, 1)) / (2 * torch.sum(fg_pos_label) + 1e-5)
dp_bg_loss = torch.sum(dp_bg_loss * torch.unsqueeze(
bg_pos_label, 1)) / (2 * torch.sum(bg_pos_label) + 1e-5)
avg_meter.add({
'loss1': pos_aff_loss.item(),
'loss2': neg_aff_loss.item(),
'loss3': dp_fg_loss.item(),
'loss4': dp_bg_loss.item()
})
total_loss = (pos_aff_loss + neg_aff_loss) / 2 + (dp_fg_loss +
dp_bg_loss) / 2
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
if (optimizer.global_step - 1) % 50 == 0:
timer.update_progress(optimizer.global_step / max_step)
print('step:%5d/%5d' % (optimizer.global_step - 1, max_step),
'loss:%.4f %.4f %.4f %.4f' %
(avg_meter.pop('loss1'), avg_meter.pop('loss2'),
avg_meter.pop('loss3'), avg_meter.pop('loss4')),
'imps:%.1f' % ((iter + 1) * args.irn_batch_size /
timer.get_stage_elapsed()),
'lr: %.4f' % (optimizer.param_groups[0]['lr']),
'etc:%s' % (timer.str_estimated_complete()),
flush=True)
else:
timer.reset_stage()
infer_dataset = dataloader.VOC12ImageDataset(args.infer_list,
voc12_root=args.voc12_root,
crop_size=args.irn_crop_size,
crop_method="top_left")
infer_data_loader = DataLoader(infer_dataset,
batch_size=args.irn_batch_size,
shuffle=False,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
model.eval()
print('Analyzing displacements mean ... ', end='')
dp_mean_list = []
with torch.no_grad():
for iter, pack in tqdm(enumerate(infer_data_loader),
total=len(infer_dataset) //
args.irn_batch_size):
img = pack['img']
if use_gpu:
img = img.cuda(non_blocking=True)
aff, dp = model(img, False)
dp_mean_list.append(torch.mean(dp, dim=(0, 2, 3)).cpu())
try:
model.module.mean_shift.running_mean = torch.mean(
torch.stack(dp_mean_list), dim=0)
except:
model.mean_shift.running_mean = torch.mean(
torch.stack(dp_mean_list), dim=0)
print('done.')
try:
state_dict = model.module.state_dict()
except:
state_dict = model.state_dict()
torch.save(state_dict, args.irn_weights_name)
if use_gpu:
torch.cuda.empty_cache()
def run(args):
assert args.train_list is not None
assert args.val_list is not None
assert args.weights_name is not None
assert args.network is not None
assert args.num_epoches is not None
assert args.network_module is not None
train_data_loader = data_loader(args, 'train')
max_step = (len(train_data_loader.dataset) //
args.batch_size) * args.num_epoches
num_classes_list = train_data_loader.dataset.get_number_classes()
model = getattr(importlib.import_module(args.network_module),
args.network)(pretrained=True,
num_classes_list=num_classes_list)
val_data_loader = data_loader(args, 'test')
params = model.trainable_parameters()
optimizer = torchutils.PolyOptimizer([{
'params': params[0],
'lr': args.learning_rate,
'weight_decay': args.weight_decay
}, {
'params': params[1],
'lr': args.learning_rate,
'weight_decay': args.weight_decay
}, {
'params': params[2],
'lr': args.learning_rate,
'weight_decay': args.weight_decay
}],
lr=args.learning_rate,
weight_decay=args.weight_decay,
max_step=max_step)
if use_gpu:
model = torch.nn.DataParallel(model).cuda()
model.train()
avg_meter = pyutils.AverageMeter()
timer = pyutils.Timer()
for ep in range(args.num_epoches):
print('Epoch %d/%d' % (ep + 1, args.num_epoches))
correct = {'make': 0.0, 'model': 0.0, 'year': 0.0}
total = 0.
for step, pack in tqdm(enumerate(train_data_loader),
total=len(train_data_loader.dataset) //
args.batch_size):
(img, path), (make_label, model_label, year_label) = pack
if use_gpu:
make_label = make_label.cuda(non_blocking=True)
model_label = model_label.cuda(non_blocking=True)
year_label = year_label.cuda(non_blocking=True)
x = model(img)
make_logit = x[0].data
model_logit = x[1].data
year_logit = x[2].data
_, make_predicted = torch.max(make_logit, 1)
_, mode_predicted = torch.max(model_logit, 1)
_, year_predicted = torch.max(year_logit, 1)
_, make_actual = torch.max(make_label.data, 1)
_, model_actual = torch.max(model_label.data, 1)
_, year_actual = torch.max(year_label.data, 1)
correct['make'] += (make_predicted == make_actual).sum()
correct['model'] += (mode_predicted == model_actual).sum()
correct['year'] += (year_predicted == year_actual).sum()
total += make_label.shape[0]
make_loss = F.multilabel_soft_margin_loss(x[0], make_label)
model_loss = F.multilabel_soft_margin_loss(x[1], model_label)
year_loss = F.multilabel_soft_margin_loss(x[2], year_label)
loss = make_loss + model_loss + year_loss
avg_meter.add({'make_loss': make_loss.item()})
avg_meter.add({'model_loss': model_loss.item()})
avg_meter.add({'year_loss': year_loss.item()})
optimizer.zero_grad()
avg_meter.add({'loss': loss.item()})
loss.backward()
optimizer.step()
if (optimizer.global_step - 1) % 100 == 0:
make_acc = 100 * correct['make'] / total
model_acc = 100 * correct['model'] / total
year_acc = 100 * correct['year'] / total
timer.update_progress(optimizer.global_step / max_step)
print(
'step:%5d/%5d' % (optimizer.global_step - 1, max_step),
'loss:%.4f' % (avg_meter.pop('loss')),
'imps:%.1f' %
((step + 1) * args.batch_size / timer.get_stage_elapsed()),
'lr: %.4f' % (optimizer.param_groups[0]['lr']),
'etc:%s' % (timer.str_estimated_complete()),
'make_acc:%s' % make_acc,
'model_acc:%s' % model_acc,
'year_acc:%s' % year_acc,
flush=True)
else:
validate(model, val_data_loader)
timer.reset_stage()
try:
state_dict = model.module.state_dict()
except:
state_dict = model.state_dict()
torch.save(state_dict, args.cam_weights_name + '.pth')
if use_gpu:
torch.cuda.empty_cache()