def main():
"""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)
model = Res_Deeplab(num_classes=args.num_classes)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
lip_dataset = LIPDataSet(args.data_dir,
'test',
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)
restore_from = args.restore_from
state_dict = model.state_dict().copy()
state_dict_old = torch.load(restore_from)
for key, nkey in zip(state_dict_old.keys(), state_dict.keys()):
if key != nkey:
# remove the 'module.' in the 'key'
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))
write_results(parsing_preds,
scales,
centers,
args.data_dir,
'test',
args.save_dir,
input_size=input_size)
def main():
"""Create the model and start the training."""
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)
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)
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)
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('iter = {} of {} completed, loss = {}'.format(
i_iter, total_iters,
loss.data.cpu().numpy()))
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'LIP_epoch_' + str(epoch) + '.pth'))
#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)
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
update_config(config, args)
print(args)
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)
model = get_cls_net(config=config,
num_classes=args.num_classes,
is_train=False)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
print('-------Load Data', args.data_dir)
if 'vehicle_parsing_dataset' in args.data_dir:
parsing_dataset = VPDataSet(args.data_dir,
args.dataset,
crop_size=input_size,
transform=transform,
list_path=args.list_path)
elif 'LIP' in args.data_dir:
parsing_dataset = LIPDataSet(args.data_dir,
args.dataset,
crop_size=input_size,
transform=transform)
num_samples = len(parsing_dataset)
valloader = data.DataLoader(parsing_dataset,
batch_size=args.batch_size * len(gpus),
shuffle=False,
pin_memory=True)
print('-------Load Weight', args.restore_from)
restore_from = args.restore_from
state_dict = model.state_dict().copy()
state_dict_old = torch.load(restore_from)
state_dict_old = state_dict_old['state_dict']
for key, nkey in zip(state_dict_old.keys(), state_dict.keys()):
if key != nkey:
# remove the 'module.' in the 'key'
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 = DataParallelModel(model)
model.eval()
model.cuda()
print('-------Start Testing...')
parsing_preds, scales, centers, time_list = valid(model, valloader,
input_size, num_samples,
len(gpus))
print('-------Saving Results', args.save_dir)
write_results(parsing_preds,
scales,
centers,
args.data_dir,
args.dataset,
args.save_dir,
input_size=input_size)
print('total time is ', sum(time_list))
print('avg time is ', sum(time_list) / len(time_list))
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, is_train=True)
# model = DataParallelModel(deeplab)
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]
print('LOAD', state_name)
else:
print('NOT LOAD', state_name)
deeplab.load_state_dict(new_params)
model = DataParallelModel(deeplab)
print('-------Load Weight', args.restore_from)
model.cuda()
criterion = CriterionAll(loss_type=args.loss)
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)
'''
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)
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."""
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
h, w = map(int, args.input_size.split(','))
input_size = [h, w]
best_f1 = 0
torch.cuda.set_device(args.local_rank)
try:
world_size = int(os.environ['WORLD_SIZE'])
distributed = world_size > 1
except:
distributed = False
world_size = 1
if distributed:
dist.init_process_group(backend=args.dist_backend,
init_method='env://')
rank = 0 if not distributed else dist.get_rank()
writer = SummaryWriter(osp.join(args.snapshot_dir,
TIMESTAMP)) if rank == 0 else None
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.type == 'Helen':
train_dataset = HelenDataSet('dataset/Helen_align_with_hair',
args.dataset,
crop_size=input_size,
transform=transform)
val_dataset = HelenDataSet('dataset/Helen_align_with_hair',
'test',
crop_size=input_size,
transform=transform)
args.num_classes = 11
elif args.type == 'LaPa':
train_dataset = LapaDataset('dataset/LaPa/origin',
args.dataset,
crop_size=input_size,
transform=transform)
val_dataset = LapaDataset('dataset/LaPa/origin',
'test',
crop_size=input_size,
transform=transform)
args.num_classes = 11
elif args.type == 'Celeb':
train_dataset = CelebAMaskHQDataSet('dataset/CelebAMask-HQ',
args.dataset,
crop_size=input_size,
transform=transform)
val_dataset = CelebAMaskHQDataSet('dataset/CelebAMask-HQ',
'test',
crop_size=input_size,
transform=transform)
args.num_classes = 19
elif args.type == 'LIP':
train_dataset = LIPDataSet('dataset/LIP',
args.dataset,
crop_size=input_size,
transform=transform)
val_dataset = LIPDataSet('dataset/LIP',
'val',
crop_size=input_size,
transform=transform)
args.num_classes = 20
if distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=2,
pin_memory=True,
drop_last=True,
sampler=train_sampler)
num_samples = len(val_dataset)
valloader = data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
drop_last=False)
cudnn.enabled = True
# cudnn related setting
cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.enabled = True
if distributed:
model = AGRNet(args.num_classes)
else:
model = AGRNet(args.num_classes, InPlaceABN)
if args.restore_from is not None:
model.load_state_dict(
torch.load(args.restore_from,
map_location='cuda:{}'.format(args.local_rank)), True)
else:
resnet_params = torch.load(
os.path.join(args.snapshot_dir, 'resnet101-imagenet.pth'))
new_params = model.state_dict().copy()
for i in resnet_params:
i_parts = i.split('.')
# print(i_parts)
if not i_parts[0] == 'fc':
new_params['.'.join(i_parts[0:])] = resnet_params[i]
model.load_state_dict(new_params)
model.cuda()
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
else:
model = SingleGPU(model)
# CriterionCrossEntropyEdgeParsing_boundary_agrnet_loss for AGRNet, CriterionCrossEntropyEdgeParsing_boundary_eagrnet_loss for EAGRNet
criterion = CriterionCrossEntropyEdgeParsing_boundary_agrnet_loss(
loss_weight=[args.l1, args.l2, args.l3, args.l4],
num_classes=args.num_classes)
criterion.cuda()
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)
for epoch in range(args.start_epoch, args.epochs):
model.train()
if distributed:
train_sampler.set_epoch(epoch)
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()
with torch.no_grad():
loss = loss.detach() * labels.shape[0]
count = labels.new_tensor([labels.shape[0]], dtype=torch.long)
if dist.is_initialized():
dist.all_reduce(count, dist.ReduceOp.SUM)
dist.all_reduce(loss, dist.ReduceOp.SUM)
loss /= count.item()
if not dist.is_initialized() or dist.get_rank() == 0:
if i_iter % 50 == 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],
args.save_num_images,
args.num_classes,
is_pred=True)
pred_edges = decode_parsing(preds[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('Edge/', edge, i_iter)
writer.add_image('Pred_edge/', pred_edge, i_iter)
print('iter = {} of {} completed, loss = {}'.format(
i_iter, total_iters,
loss.data.cpu().numpy()))
if not dist.is_initialized() or dist.get_rank() == 0:
save_path = os.path.join(args.data_dir, TIMESTAMP)
if not os.path.exists(save_path):
os.makedirs(save_path)
parsing_preds, scales, centers = valid(
model, valloader, input_size, num_samples,
osp.join(args.snapshot_dir, save_path))
mIoU, f1 = compute_mean_ioU(parsing_preds,
scales,
centers,
args.num_classes,
val_dataset,
input_size,
'test',
True,
type=args.type)
if f1['mean'] > best_f1:
torch.save(model.module.state_dict(),
osp.join(args.snapshot_dir, TIMESTAMP, 'best.pth'))
best_f1 = f1['mean']
print(mIoU)
print(f1)
writer.add_scalars('mIoU', mIoU, epoch)
writer.add_scalars('f1', f1, epoch)
if epoch % args.test_fre == 0:
torch.save(
model.module.state_dict(),
osp.join(args.snapshot_dir, TIMESTAMP,
'epoch_' + str(epoch) + '.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
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.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()
state_dict_pretrain = saved_state_dict['state_dict']
for i 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[i]
else:
print('NOT LOAD', state_name)
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)
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)
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)
# 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 % 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')