def main():
"""Create the model and start the training."""
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
logger = util.set_logger(args.snapshot_dir, args.log_file, args.debug)
logger.info('start with arguments %s', args)
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
lscale, hscale = map(float, args.train_scale.split(','))
train_scale = (lscale, hscale)
cudnn.enabled = True
# Create network.
model = Res_Deeplab(num_classes=args.num_classes)
#saved_state_dict = torch.load(args.restore_from)
#new_params = model.state_dict().copy()
#for i in saved_state_dict:
# #Scale.layer5.conv2d_list.3.weight
# i_parts = i.split('.')
# # print i_parts
# if not args.num_classes == 21 or not i_parts[1]=='layer5':
# new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
model_urls = {
'resnet18':
'https://download.pytorch.org/models/resnet18-5c106cde.pth',
'resnet34':
'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
'resnet50':
'https://download.pytorch.org/models/resnet50-19c8e357.pth',
'resnet101':
'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
'resnet152':
'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}
saved_state_dict = torch.utils.model_zoo.load_url(model_urls['resnet101'])
# coco pretrained parameters:
# saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
#Scale.layer5.conv2d_list.3.weight
i_parts = str(i).split('.')
# print i_parts
if not i_parts[0] == 'fc':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
model.load_state_dict(new_params)
#model.float()
model.eval() # use_global_stats = True
#model.train()
device = torch.device("cuda:" + str(args.gpu))
model.to(device)
cudnn.benchmark = True
trainloader = data.DataLoader(GTA5DataSet(args.data_dir,
args.data_list,
max_iters=args.num_steps *
args.batch_size,
crop_size=input_size,
train_scale=train_scale,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN,
std=IMG_STD),
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=args.pin_memory)
optimizer = optim.SGD([{
'params': get_1x_lr_params_NOscale(model),
'lr': args.learning_rate
}, {
'params': get_10x_lr_params(model),
'lr': 10 * args.learning_rate
}],
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
interp = nn.Upsample(size=input_size, mode='bilinear', align_corners=True)
for i_iter, batch in enumerate(trainloader):
images, labels, _, _ = batch
images = images.to(device)
labels = labels.long().to(device)
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
pred = interp(model(images))
loss = loss_calc(pred, labels)
loss.backward()
optimizer.step()
# print('iter = ', i_iter, 'of', args.num_steps,'completed, loss = ', loss.data.cpu().numpy())
logger.info('iter = {} of {} completed, loss = {:.4f}'.format(
i_iter, args.num_steps,
loss.data.cpu().numpy()))
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'VOC12_scenes_' + str(args.num_steps) + '.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'VOC12_scenes_' + str(i_iter) + '.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
"""Create the model and start the evaluation process."""
device = torch.device("cuda:" + str(args.gpu))
if not os.path.exists(args.save):
os.makedirs(args.save)
logger = util.set_logger(args.save, args.log_file, args.debug)
logger.info('start with arguments %s', args)
x_num = 0
with open(args.data_list) as f:
for _ in f.readlines():
x_num = x_num + 1
sys.path.insert(0, 'dataset/helpers')
if args.data_src == 'gta' or args.data_src == 'cityscapes':
from labels import id2label, trainId2label
elif args.data_src == 'synthia':
from labels_cityscapes_synthia import id2label, trainId2label
#
label_2_id = 255 * np.ones((256, ))
for l in id2label:
if l in (-1, 255):
continue
label_2_id[l] = id2label[l].trainId
id_2_label = np.array(
[trainId2label[_].id for _ in trainId2label if _ not in (-1, 255)])
valid_labels = sorted(set(id_2_label.ravel()))
scorer = ScoreUpdater(valid_labels, args.num_classes, x_num, logger)
scorer.reset()
if args.model == 'DeeplabRes':
model = Res_Deeplab(num_classes=args.num_classes)
# elif args.model == 'DeeplabVGG':
# model = DeeplabVGG(num_classes=args.num_classes)
# if args.restore_from == RESTORE_FROM:
# args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
# Scale.layer5.conv2d_list.3.weight
i_parts = str(i).split('.')
# print i_parts
if not i_parts[0] == 'fc':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
else:
loc = "cuda:" + str(args.gpu)
saved_state_dict = torch.load(args.restore_from, map_location=loc)
new_params = saved_state_dict.copy()
model.load_state_dict(new_params)
#model.train()
model.eval()
model.to(device)
testloader = data.DataLoader(GTA5TestDataSet(args.data_dir,
args.data_list,
test_scale=1.0,
test_size=(1024, 512),
mean=IMG_MEAN,
std=IMG_STD,
scale=False,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
test_scales = [float(_) for _ in str(args.test_scale).split(',')]
h, w = map(int, args.test_image_size.split(','))
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(h, w), mode='bilinear', align_corners=True)
else:
interp = nn.Upsample(size=(h, w), mode='bilinear')
test_image_size = (h, w)
mean_rgb = IMG_MEAN[::-1].copy()
std_rgb = IMG_STD[::-1].copy()
with torch.no_grad():
for index, batch in enumerate(testloader):
image, label, _, name = batch
img = image.clone()
num_scales = len(test_scales)
# output_dict = {k: [] for k in range(num_scales)}
for scale_idx in range(num_scales):
if version.parse(torch.__version__) > version.parse('0.4.0'):
image = F.interpolate(image,
scale_factor=test_scales[scale_idx],
mode='bilinear',
align_corners=True)
else:
test_size = (int(h * test_scales[scale_idx]),
int(w * test_scales[scale_idx]))
interp_tmp = nn.Upsample(size=test_size,
mode='bilinear',
align_corners=True)
image = interp_tmp(img)
if args.model == 'DeeplabRes':
output2 = model(image.to(device))
coutput = interp(output2).cpu().data[0].numpy()
if args.test_flipping:
output2 = model(
torch.from_numpy(
image.numpy()[:, :, :, ::-1].copy()).to(device))
coutput = 0.5 * (
coutput +
interp(output2).cpu().data[0].numpy()[:, :, ::-1])
if scale_idx == 0:
output = coutput.copy()
else:
output += coutput
output = output / num_scales
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
pred_label = output.copy()
label = label_2_id[np.asarray(label.numpy(), dtype=np.uint8)]
scorer.update(pred_label.flatten(), label.flatten(), index)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' %
(args.save, name.split('.')[0]))
