def main():
"""Create the model and start the training."""
print(args)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
writer = SummaryWriter(args.snapshot_dir)
gpus = [int(i) for i in args.gpu.split(',')]
if not args.gpu == 'None':
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
h, w = map(int, args.input_size.split(','))
input_size = [h, w]
# cudnn related setting
cudnn.enabled = True
cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.enabled = True
deeplab = get_cls_net(config=config,
num_classes=args.num_classes,
is_train=True)
print('-------Load Weight', args.restore_from)
saved_state_dict = torch.load(args.restore_from)
if args.start_epoch > 0:
model = DataParallelModel(deeplab)
model.load_state_dict(saved_state_dict['state_dict'])
else:
new_params = deeplab.state_dict().copy()
state_dict_pretrain = saved_state_dict
for state_name in state_dict_pretrain:
if state_name in new_params:
new_params[state_name] = state_dict_pretrain[state_name]
else:
print('NOT LOAD', state_name)
deeplab.load_state_dict(new_params)
model = DataParallelModel(deeplab)
print('-------Load Weight Finish', args.restore_from)
model.cuda()
criterion = CriterionAll()
criterion = DataParallelCriterion(criterion)
criterion.cuda()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
print("-------Loading data...")
if 'vehicle_parsing_dataset' in args.data_dir:
parsing_dataset = VPDataSet(args.data_dir,
args.dataset,
crop_size=input_size,
transform=transform)
elif 'LIP' in args.data_dir:
parsing_dataset = LIPDataSet(args.data_dir,
args.dataset,
crop_size=input_size,
transform=transform)
print("Data dir : ", args.data_dir)
print("Dataset : ", args.dataset)
trainloader = data.DataLoader(parsing_dataset,
batch_size=args.batch_size * len(gpus),
shuffle=True,
num_workers=8,
pin_memory=True)
optimizer = optim.SGD(model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.start_epoch > 0:
optimizer.load_state_dict(saved_state_dict['optimizer'])
print('-------Load Optimizer', args.restore_from)
print("-------Start training...")
total_iters = args.epochs * len(trainloader)
for epoch in range(args.start_epoch, args.epochs):
model.train()
for i_iter, batch in enumerate(trainloader):
i_iter += len(trainloader) * epoch
lr = adjust_learning_rate(optimizer, i_iter, total_iters)
images, labels, _ = batch
labels = labels.long().cuda(non_blocking=True)
preds = model(images)
loss = criterion(preds, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i_iter % 100 == 0:
writer.add_scalar('learning_rate', lr, i_iter)
writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter)
print(
f'epoch = {epoch}, iter = {i_iter}/{total_iters}, lr={lr:.6f}, loss = {loss.data.cpu().numpy():.6f}'
)
if (epoch + 1) % args.save_step == 0 or epoch == args.epochs:
time.sleep(10)
print("-------Saving checkpoint...")
save_checkpoint(model, epoch, optimizer)
time.sleep(10)
save_checkpoint(model, epoch, optimizer)
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
"""Create the model and start the training."""
print(args)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
writer = SummaryWriter(args.snapshot_dir)
gpus = [int(i) for i in args.gpu.split(',')]
if not args.gpu == 'None':
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
h, w = map(int, args.input_size.split(','))
input_size = [h, w]
cudnn.enabled = True
# cudnn related setting
cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.enabled = True
deeplab = get_cls_net(config=config,
num_classes=args.num_classes,
is_train=True)
model = DataParallelModel(deeplab)
saved_state_dict = torch.load(args.restore_from)
if args.start_epoch > 0:
model = DataParallelModel(deeplab)
model.load_state_dict(saved_state_dict['state_dict'])
else:
new_params = model.state_dict().copy()
state_dict_pretrain = saved_state_dict['state_dict']
for state_name in state_dict_pretrain:
if state_name in new_params:
new_params[state_name] = state_dict_pretrain[state_name]
#print ('LOAD',state_name)
else:
print('NOT LOAD', state_name)
model.load_state_dict(new_params)
print('-------Load Weight', args.restore_from)
model.cuda()
criterion = CriterionAll2()
criterion = DataParallelCriterion(criterion)
criterion.cuda()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
trainloader = data.DataLoader(LIPDataSet(args.data_dir,
args.dataset,
crop_size=input_size,
transform=transform),
batch_size=args.batch_size * len(gpus),
shuffle=True,
num_workers=4,
pin_memory=True)
num_samples = 5000
'''
list_map = []
for part in deeplab.path_list:
list_map = list_map + list(map(id, part.parameters()))
base_params = filter(lambda p: id(p) not in list_map,
deeplab.parameters())
params_list = []
params_list.append({'params': base_params, 'lr':args.learning_rate*0.1})
for part in deeplab.path_list:
params_list.append({'params': part.parameters()})
print ('len(params_list)',len(params_list))
'''
list_map = []
for part in deeplab.path_list:
list_map = list_map + list(map(id, part.parameters()))
base_params = filter(lambda p: id(p) not in list_map, deeplab.parameters())
params_list = []
params_list.append({'params': base_params, 'lr': 1e-6})
for part in deeplab.path_list:
params_list.append({'params': part.parameters()})
print('len(params_list)', len(params_list))
optimizer = torch.optim.SGD(params_list,
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.start_epoch > 0:
optimizer.load_state_dict(saved_state_dict['optimizer'])
print('========Load Optimizer', args.restore_from)
optimizer.zero_grad()
total_iters = args.epochs * len(trainloader)
for epoch in range(args.start_epoch, args.epochs):
model.train()
for i_iter, batch in enumerate(trainloader):
i_iter += len(trainloader) * epoch
#lr = adjust_learning_rate(optimizer, i_iter, total_iters)
lr = adjust_learning_rate_parsing(optimizer, epoch)
images, labels, _ = batch
labels = labels.long().cuda(non_blocking=True)
preds = model(images)
loss = criterion(preds, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i_iter % 100 == 0:
writer.add_scalar('learning_rate', lr, i_iter)
writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter)
print('epoch = {}, iter = {} of {} completed,lr={}, loss = {}'.
format(epoch, i_iter, total_iters, lr,
loss.data.cpu().numpy()))
if epoch % 2 == 0 or epoch == args.epochs:
time.sleep(10)
save_checkpoint(model, epoch, optimizer)
# parsing_preds, scales, centers = valid(model, valloader, input_size, num_samples, len(gpus))
# mIoU = compute_mean_ioU(parsing_preds, scales, centers, args.num_classes, args.data_dir, input_size)
# print(mIoU)
# writer.add_scalars('mIoU', mIoU, epoch)
time.sleep(10)
save_checkpoint(model, epoch, optimizer)
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
"""Create the model and start the training."""
print (args)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
writer = SummaryWriter(args.snapshot_dir)
gpus = [int(i) for i in args.gpu.split(',')]
if not args.gpu == 'None':
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
h, w = map(int, args.input_size.split(','))
input_size = [h, w]
cudnn.enabled = True
# cudnn related setting
cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.enabled = True
deeplab = Res_Deeplab(num_classes=args.num_classes)
# dump_input = torch.rand((args.batch_size, 3, input_size[0], input_size[1]))
# writer.add_graph(deeplab.cuda(), dump_input.cuda(), verbose=False)
saved_state_dict = torch.load(args.restore_from)
if args.start_epoch >0:
model = DataParallelModel(deeplab)
model.load_state_dict(saved_state_dict['state_dict'])
else:
new_params = deeplab.state_dict().copy()
for i in saved_state_dict:
i_parts = i.split('.')
# print(i_parts)
if not i_parts[0] == 'fc':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
deeplab.load_state_dict(new_params)
model = DataParallelModel(deeplab)
print ('-------Load Weight',args.restore_from)
model.cuda()
criterion = CriterionAll()
criterion = DataParallelCriterion(criterion)
criterion.cuda()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
trainloader = data.DataLoader(LIPDataSet(args.data_dir, args.dataset, crop_size=input_size, transform=transform),
batch_size=args.batch_size * len(gpus), shuffle=True, num_workers=4,
pin_memory=True)
lip_dataset = LIPDataSet(args.data_dir, 'val', crop_size=input_size, transform=transform)
num_samples = len(lip_dataset)
valloader = data.DataLoader(lip_dataset, batch_size=args.batch_size * len(gpus),
shuffle=False, pin_memory=True)
optimizer = optim.SGD(
model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay
)
if args.start_epoch > 0:
optimizer.load_state_dict(saved_state_dict['optimizer'])
print ('========Load Optimizer',args.restore_from)
optimizer.zero_grad()
total_iters = args.epochs * len(trainloader)
for epoch in range(args.start_epoch, args.epochs):
model.train()
for i_iter, batch in enumerate(trainloader):
i_iter += len(trainloader) * epoch
lr = adjust_learning_rate(optimizer, i_iter, total_iters)
images, labels, edges, _ = batch
labels = labels.long().cuda(non_blocking=True)
edges = edges.long().cuda(non_blocking=True)
preds = model(images)
loss = criterion(preds, [labels, edges])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i_iter % 100 == 0:
writer.add_scalar('learning_rate', lr, i_iter)
writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter)
# if i_iter % 500 == 0:
# images_inv = inv_preprocess(images, args.save_num_images)
# labels_colors = decode_parsing(labels, args.save_num_images, args.num_classes, is_pred=False)
# edges_colors = decode_parsing(edges, args.save_num_images, 2, is_pred=False)
# if isinstance(preds, list):
# preds = preds[0]
# preds_colors = decode_parsing(preds[0][-1], args.save_num_images, args.num_classes, is_pred=True)
# pred_edges = decode_parsing(preds[1][-1], args.save_num_images, 2, is_pred=True)
# img = vutils.make_grid(images_inv, normalize=False, scale_each=True)
# lab = vutils.make_grid(labels_colors, normalize=False, scale_each=True)
# pred = vutils.make_grid(preds_colors, normalize=False, scale_each=True)
# edge = vutils.make_grid(edges_colors, normalize=False, scale_each=True)
# pred_edge = vutils.make_grid(pred_edges, normalize=False, scale_each=True)
# writer.add_image('Images/', img, i_iter)
# writer.add_image('Labels/', lab, i_iter)
# writer.add_image('Preds/', pred, i_iter)
# writer.add_image('Edges/', edge, i_iter)
# writer.add_image('PredEdges/', pred_edge, i_iter)
print('epoch = {}, iter = {} of {} completed,lr={}, loss = {}'.format(epoch, i_iter, total_iters,lr, loss.data.cpu().numpy()))
if epoch%5 == 0 or epoch==args.epochs:
time.sleep(10)
save_checkpoint(model,epoch,optimizer)
# parsing_preds, scales, centers = valid(model, valloader, input_size, num_samples, len(gpus))
# mIoU = compute_mean_ioU(parsing_preds, scales, centers, args.num_classes, args.data_dir, input_size)
# print(mIoU)
# writer.add_scalars('mIoU', mIoU, epoch)
time.sleep(10)
save_checkpoint(model,epoch,optimizer)
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
torch.multiprocessing.set_start_method("spawn", force=True)
"""Create the model and start the evaluation process."""
args = get_arguments()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
gpus = [int(i) for i in args.gpu.split(',')]
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
deeplab = CorrPM_Model(args.num_classes, args.num_points)
if len(gpus) > 1:
model = DataParallelModel(deeplab)
else:
model = deeplab
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
if args.data_name == 'lip':
lip_dataset = LIPDataSet(args.data_dir,
VAL_POSE_ANNO_FILE,
args.dataset,
crop_size=input_size,
transform=transform)
num_samples = len(lip_dataset)
valloader = data.DataLoader(lip_dataset,
batch_size=args.batch_size * len(gpus),
shuffle=False,
num_workers=4,
pin_memory=True)
restore_from = args.restore_from
state_dict = model.state_dict().copy()
state_dict_old = torch.load(restore_from)
for key in state_dict.keys():
if key not in state_dict_old.keys():
print(key)
for key, nkey in zip(state_dict_old.keys(), state_dict.keys()):
if key != nkey:
state_dict[key[7:]] = deepcopy(state_dict_old[key])
else:
state_dict[key] = deepcopy(state_dict_old[key])
model.load_state_dict(state_dict)
model.eval()
model.cuda()
parsing_preds, scales, centers = valid(model, valloader, input_size,
num_samples, len(gpus))
mIoU = compute_mean_ioU(parsing_preds, scales, centers, args.num_classes,
args.data_dir, input_size, args.dataset)
print(mIoU)
end = datetime.datetime.now()
print(end - start, 'seconds')
print(end)
def main():
"""Create the model and start the training."""
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
timestramp = args.date
writer = SummaryWriter(os.path.join(args.snapshot_dir, timestramp))
gpus = [int(i) for i in args.gpu.split(',')]
if not args.gpu == 'None':
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
h, w = map(int, args.input_size.split(','))
input_size = [h, w]
cudnn.enabled = True
# cudnn related setting
cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.enabled = True
deeplab = Res_Deeplab(num_classes=args.num_classes)
# dump_input = torch.rand((args.batch_size, 3, input_size[0], input_size[1]))
# writer.add_graph(deeplab.cuda(), dump_input.cuda(), verbose=False)
model = DataParallelModel(deeplab)
if args.resume:
# when restore form the same network, it is useful here
checkpoint = torch.load(args.restore_from)
model.load_state_dict(checkpoint['net'])
args.start_epoch = checkpoint['epoch']
else:
saved_state_dict = torch.load(args.restore_from)
new_params = deeplab.state_dict().copy()
for i in saved_state_dict:
i_parts = i.split('.')
if not i_parts[0] == 'fc':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
deeplab.load_state_dict(new_params)
model.cuda()
criterion = CriterionAll()
criterion = DataParallelCriterion(criterion)
criterion.cuda()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
trainloader = data.DataLoader(LIPDataSet(args.data_dir,
args.dataset,
crop_size=input_size,
transform=transform),
batch_size=args.batch_size * len(gpus),
shuffle=True,
num_workers=2,
pin_memory=True)
lip_dataset = LIPDataSet(args.data_dir,
'val',
crop_size=input_size,
transform=transform)
num_samples = len(lip_dataset)
valloader = data.DataLoader(lip_dataset,
batch_size=args.batch_size * len(gpus),
shuffle=False,
pin_memory=True)
optimizer = optim.SGD(model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
total_iters = args.epochs * len(trainloader)
log = Logger(os.path.join(args.log_dir, '{}_train.log'.format(timestramp)),
level='debug')
for epoch in range(args.start_epoch, args.epochs):
model.train()
#lr = adjust_learning_rate_pose(optimizer, epoch)
for i_iter, batch in enumerate(trainloader):
i_iter += len(trainloader) * epoch
lr = adjust_learning_rate(optimizer, i_iter, total_iters)
images, labels, r1, r2, r3, r4, l0, l1, l2, l3, l4, l5, _ = batch
labels = labels.long().cuda(non_blocking=True)
r1 = r1.long().cuda(non_blocking=True)
r2 = r2.long().cuda(non_blocking=True)
r3 = r3.long().cuda(non_blocking=True)
r4 = r4.long().cuda(non_blocking=True)
l0 = l0.long().cuda(non_blocking=True)
l1 = l1.long().cuda(non_blocking=True)
l2 = l2.long().cuda(non_blocking=True)
l3 = l3.long().cuda(non_blocking=True)
l4 = l4.long().cuda(non_blocking=True)
l5 = l5.long().cuda(non_blocking=True)
preds = model(images)
loss = criterion(
preds, [[labels], [r1, r2, r3, r4], [l0, l1, l2, l3, l4, l5]])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i_iter % 100 == 0:
writer.add_scalar('learning_rate', lr, i_iter)
writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter)
if i_iter % 10 == 0:
log.logger.info(
'epoch = {} iter = {} of {} completed, lr = {}, loss = {}'.
format(epoch, i_iter, total_iters, lr,
loss.data.cpu().numpy()))
parsing_preds, scales, centers = valid(model, valloader, input_size,
num_samples, len(gpus))
mIoU = compute_mean_ioU(parsing_preds, scales, centers,
args.num_classes, args.data_dir, input_size)
log.logger.info('epoch = {}'.format(epoch))
log.logger.info(str(mIoU))
writer.add_scalars('mIoU', mIoU, epoch)
# save the model snapshot
state = {"net": model.module.state_dict(), "epoch": epoch}
torch.save(
state,
osp.join(args.snapshot_dir, timestramp,
'LIP_epoch_' + str(epoch) + '.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
"""Create the model and start the training."""
print(args)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
writer = SummaryWriter(args.snapshot_dir)
gpus = [int(i) for i in args.gpu.split(',')]
if not args.gpu == 'None':
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
h, w = map(int, args.input_size.split(','))
input_size = [h, w]
# cudnn related setting
cudnn.enabled = True
cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.enabled = True
deeplab = get_resnet101_asp_oc_dsn(num_classes=args.num_classes)
# dump_input = torch.rand((args.batch_size, 3, input_size[0], input_size[1]))
# writer.add_graph(deeplab.cuda(), dump_input.cuda(), verbose=False)
saved_state_dict = torch.load(args.restore_from)
if args.start_epoch > 0:
model = DataParallelModel(deeplab)
#model = torch.nn.parallel.DistributedDataParallel(deeplab)
model.load_state_dict(saved_state_dict['state_dict'])
else:
new_params = deeplab.state_dict().copy()
state_dict_pretrain = saved_state_dict #['state_dict']
for state_name in state_dict_pretrain:
# splits = i.split('.')
# state_name = '.'.join(splits[1:])
if state_name in new_params:
new_params[state_name] = state_dict_pretrain[state_name]
else:
print('NOT LOAD', state_name)
deeplab.load_state_dict(new_params)
model = DataParallelModel(deeplab)
#model = torch.nn.parallel.DistributedDataParallel(deeplab)
print('-------Load Weight', args.restore_from)
model.cuda()
criterion = LovaszSoftmaxDSN(input_size)
print('LOSS1: LovaszSoftmaxDSN')
criterion = DataParallelCriterion(criterion)
criterion.cuda()
criterion_softmax = CriterionDSN()
print('LOSS2: CriterionDSN')
criterion_softmax = DataParallelCriterion(criterion_softmax)
criterion_softmax.cuda()
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
trainloader = data.DataLoader(LIPDataSet(args.data_dir,
args.dataset,
crop_size=input_size,
transform=transform,
list_path=args.list_path),
batch_size=args.batch_size * len(gpus),
shuffle=True,
num_workers=4,
pin_memory=True)
num_samples = 5000
'''
list_map = []
for part in deeplab.path_list:
list_map = list_map + list(map(id, part.parameters()))
base_params = filter(lambda p: id(p) not in list_map,
deeplab.parameters())
params_list = []
params_list.append({'params': base_params, 'lr':args.learning_rate*0.1})
for part in deeplab.path_list:
params_list.append({'params': part.parameters()})
print ('len(params_list)',len(params_list))
'''
optimizer = optim.SGD(model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.start_epoch > 0:
optimizer.load_state_dict(saved_state_dict['optimizer'])
print('========Load Optimizer', args.restore_from)
total_iters = args.epochs * len(trainloader)
for epoch in range(args.start_epoch, args.epochs):
model.train()
for i_iter, batch in enumerate(trainloader):
i_iter += len(trainloader) * epoch
lr = adjust_learning_rate(optimizer, i_iter, total_iters)
images, labels, _ = batch
labels = labels.long().cuda(non_blocking=True)
preds = model(images)
loss1 = criterion(preds, labels)
loss2 = criterion_softmax(preds, labels)
loss = loss1 + loss2
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i_iter % 100 == 0:
writer.add_scalar('learning_rate', lr, i_iter)
writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter)
# if i_iter % 500 == 0:
# images_inv = inv_preprocess(images, args.save_num_images)
# labels_colors = decode_parsing(labels, args.save_num_images, args.num_classes, is_pred=False)
# edges_colors = decode_parsing(edges, args.save_num_images, 2, is_pred=False)
# if isinstance(preds, list):
# preds = preds[0]
# preds_colors = decode_parsing(preds[0][-1], args.save_num_images, args.num_classes, is_pred=True)
# pred_edges = decode_parsing(preds[1][-1], args.save_num_images, 2, is_pred=True)
# img = vutils.make_grid(images_inv, normalize=False, scale_each=True)
# lab = vutils.make_grid(labels_colors, normalize=False, scale_each=True)
# pred = vutils.make_grid(preds_colors, normalize=False, scale_each=True)
# edge = vutils.make_grid(edges_colors, normalize=False, scale_each=True)
# pred_edge = vutils.make_grid(pred_edges, normalize=False, scale_each=True)
# writer.add_image('Images/', img, i_iter)
# writer.add_image('Labels/', lab, i_iter)
# writer.add_image('Preds/', pred, i_iter)
# writer.add_image('Edges/', edge, i_iter)
# writer.add_image('PredEdges/', pred_edge, i_iter)
print(
'epoch = {}, iter = {} of {} completed,lr={:.4f}, loss = {:.4f}, IoU_loss = {:.4f}, BCE_loss = {:.4f}'
.format(epoch, i_iter, total_iters, lr,
loss.data.cpu().numpy(),
loss1.data.cpu().numpy(),
loss2.data.cpu().numpy()))
if epoch % args.save_step == 0 or epoch == args.epochs:
time.sleep(10)
save_checkpoint(model, epoch, optimizer)
# parsing_preds, scales, centers = valid(model, valloader, input_size, num_samples, len(gpus))
# mIoU = compute_mean_ioU(parsing_preds, scales, centers, args.num_classes, args.data_dir, input_size)
# print(mIoU)
# writer.add_scalars('mIoU', mIoU, epoch)
time.sleep(10)
save_checkpoint(model, epoch, optimizer)
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
args = get_arguments()
print(args)
start_epoch = 0
cycle_n = 0
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
with open(os.path.join(args.log_dir, 'args.json'), 'w') as opt_file:
json.dump(vars(args), opt_file)
gpus = [int(i) for i in args.gpu.split(',')]
if not args.gpu == 'None':
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
input_size = list(map(int, args.input_size.split(',')))
cudnn.enabled = True
cudnn.benchmark = True
# Model Initialization
AugmentCE2P = networks.init_model(args.arch,
num_classes=args.num_classes,
pretrained=args.imagenet_pretrain)
model = DataParallelModel(AugmentCE2P)
model.cuda()
IMAGE_MEAN = AugmentCE2P.mean
IMAGE_STD = AugmentCE2P.std
INPUT_SPACE = AugmentCE2P.input_space
print('image mean: {}'.format(IMAGE_MEAN))
print('image std: {}'.format(IMAGE_STD))
print('input space:{}'.format(INPUT_SPACE))
restore_from = args.model_restore
if os.path.exists(restore_from):
print('Resume training from {}'.format(restore_from))
checkpoint = torch.load(restore_from)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
SCHP_AugmentCE2P = networks.init_model(args.arch,
num_classes=args.num_classes,
pretrained=args.imagenet_pretrain)
schp_model = DataParallelModel(SCHP_AugmentCE2P)
schp_model.cuda()
if os.path.exists(args.schp_restore):
print('Resuming schp checkpoint from {}'.format(args.schp_restore))
schp_checkpoint = torch.load(args.schp_restore)
schp_model_state_dict = schp_checkpoint['state_dict']
cycle_n = schp_checkpoint['cycle_n']
schp_model.load_state_dict(schp_model_state_dict)
# Loss Function
criterion = CriterionAll(lambda_1=args.lambda_s,
lambda_2=args.lambda_e,
lambda_3=args.lambda_c,
num_classes=args.num_classes)
criterion = DataParallelCriterion(criterion)
criterion.cuda()
# Data Loader
if INPUT_SPACE == 'BGR':
print('BGR Transformation')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=IMAGE_MEAN, std=IMAGE_STD),
])
elif INPUT_SPACE == 'RGB':
print('RGB Transformation')
transform = transforms.Compose([
transforms.ToTensor(),
BGR2RGB_transform(),
transforms.Normalize(mean=IMAGE_MEAN, std=IMAGE_STD),
])
train_dataset = LIPDataSet(args.data_dir,
args.split_name,
crop_size=input_size,
transform=transform)
train_loader = data.DataLoader(train_dataset,
batch_size=args.batch_size * len(gpus),
num_workers=16,
shuffle=True,
pin_memory=True,
drop_last=True)
print('Total training samples: {}'.format(len(train_dataset)))
# Optimizer Initialization
optimizer = optim.SGD(model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
lr_scheduler = SGDRScheduler(optimizer,
total_epoch=args.epochs,
eta_min=args.learning_rate / 100,
warmup_epoch=10,
start_cyclical=args.schp_start,
cyclical_base_lr=args.learning_rate / 2,
cyclical_epoch=args.cycle_epochs)
total_iters = args.epochs * len(train_loader)
start = timeit.default_timer()
for epoch in range(start_epoch, args.epochs):
lr_scheduler.step(epoch=epoch)
lr = lr_scheduler.get_lr()[0]
model.train()
for i_iter, batch in enumerate(train_loader):
i_iter += len(train_loader) * epoch
images, labels, _ = batch
labels = labels.cuda(non_blocking=True)
edges = generate_edge_tensor(labels)
labels = labels.type(torch.cuda.LongTensor)
edges = edges.type(torch.cuda.LongTensor)
preds = model(images)
# Online Self Correction Cycle with Label Refinement
if cycle_n >= 1:
with torch.no_grad():
soft_preds = schp_model(images)
soft_parsing = []
soft_edge = []
for soft_pred in soft_preds:
soft_parsing.append(soft_pred[0][-1])
soft_edge.append(soft_pred[1][-1])
soft_preds = torch.cat(soft_parsing, dim=0)
soft_edges = torch.cat(soft_edge, dim=0)
else:
soft_preds = None
soft_edges = None
loss = criterion(preds, [labels, edges, soft_preds, soft_edges],
cycle_n)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i_iter % 100 == 0:
print('iter = {} of {} completed, lr = {}, loss = {}'.format(
i_iter, total_iters, lr,
loss.data.cpu().numpy()))
if (epoch + 1) % (args.eval_epochs) == 0:
schp.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
},
False,
args.log_dir,
filename='checkpoint_{}.pth.tar'.format(epoch + 1))
# Self Correction Cycle with Model Aggregation
if (epoch + 1) >= args.schp_start and (
epoch + 1 - args.schp_start) % args.cycle_epochs == 0:
print('Self-correction cycle number {}'.format(cycle_n))
schp.moving_average(schp_model, model, 1.0 / (cycle_n + 1))
cycle_n += 1
schp.bn_re_estimate(train_loader, schp_model)
schp.save_schp_checkpoint(
{
'state_dict': schp_model.state_dict(),
'cycle_n': cycle_n,
},
False,
args.log_dir,
filename='schp_{}_checkpoint.pth.tar'.format(cycle_n))
torch.cuda.empty_cache()
end = timeit.default_timer()
print('epoch = {} of {} completed using {} s'.format(
epoch, args.epochs, (end - start) / (epoch - start_epoch + 1)))
end = timeit.default_timer()
print('Training Finished in {} seconds'.format(end - start))