def main():
"""Create the model and start the training."""
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
gpu = args.gpu
# Create network
if args.model == 'DeepLab':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
# ====the lines below comment by ZQ====
# 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 == 19 or not i_parts[1] == 'layer5':
# new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# # print i_parts
# model.load_state_dict(new_params)
# ====end====
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# init D
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(MixedDataSet(
args.data_dir,
args.data_list,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer_D1 = optim.Adam(model_D1.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
# optimizer_D1.zero_grad()
#
# optimizer_D2 = optim.Adam(model_D2.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
# optimizer_D2.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
# interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear')
# labels for adversarial training
# source_label = 0
# target_label = 1
loss2_sum_per_epoch = 0
loss2_per_epoch = 0
epoch = 0
loss2_epoch = ''
lbl_list = open(args.data_list, 'r')
lbl_num = len(lbl_list.readlines()) / 2
mIoUs = []
i_iters = []
syn_mIoUs = []
for i_iter in range(args.num_steps):
loss_seg_value1 = 0
# loss_adv_target_value1 = 0
# loss_D_value1 = 0
loss_seg_value2 = 0
# loss_adv_target_value2 = 0
# loss_D_value2 = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
for sub_i in range(args.iter_size):
# train G
# train with source
_, batch = trainloader_iter.__next__()
images, labels, _, img_nam = batch
# print('%s\n' % img_nam[0])
images = Variable(images).cuda(args.gpu)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
loss_seg1 = loss_calc(pred1, labels, args.gpu)
loss_seg2 = loss_calc(pred2, labels, args.gpu)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value1 += loss_seg1.data.cpu().numpy() / args.iter_size
loss_seg_value2 += loss_seg2.data.cpu().numpy() / args.iter_size
optimizer.step()
# print('exp = {}'.format(args.snapshot_dir))
print('iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f}'.
format(i_iter, args.num_steps, loss_seg_value1, loss_seg_value2))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + 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,
'GTA5_' + str(i_iter + strat_snap_iter) + '.pth'))
show_pred_sv_dir = pre_sv_dir.format(i_iter)
mIoU, syn_mIoU = show_val(model.state_dict(), show_pred_sv_dir)
mIoUs.append(str(round(np.nanmean(mIoU) * 100, 2)))
i_iters.append(i_iter + strat_snap_iter)
syn_mIoUs.append(str(round(np.nanmean(syn_mIoU) * 100, 2)))
print_i = 0
for miou in mIoUs:
print('i{0}: {1} {2}'.format(i_iters[print_i], miou,
syn_mIoUs[print_i]))
print_i = print_i + 1
def main():
"""Create the model and start the training."""
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
h, w = map(int, args.input_size_target.split(','))
input_size_target = (h, w)
cudnn.enabled = True
gpu = args.gpu
# Create network
if args.model == 'DeepLab':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
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 == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# print i_parts
model.load_state_dict(new_params)
# model.load_state_dict(saved_state_dict)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes)
model_D2 = FCDiscriminator(num_classes=args.num_classes)
# # load discriminator params
# saved_state_dict_D1 = torch.load(D1_RESTORE_FROM)
# saved_state_dict_D2 = torch.load(D2_RESTORE_FROM)
# model_D1.load_state_dict(saved_state_dict_D1)
# model_D2.load_state_dict(saved_state_dict_D2)
model_D1.train()
model_D1.cuda(args.gpu)
model_D2.train()
model_D2.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(synthiaDataSet(
args.data_dir,
args.data_list,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False,
mirror=args.random_mirror,
mean=CITY_IMG_MEAN,
set=args.set),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer_D1 = optim.Adam(model_D1.parameters(),
lr=args.learning_rate_D,
betas=(0.7, 0.99))
optimizer_D2 = optim.Adam(model_D2.parameters(),
lr=args.learning_rate_D,
betas=(0.7, 0.99))
#BYZQ
# opti_state_dict = torch.load(OPTI_RESTORE_FROM)
# opti_state_dict_d1 = torch.load(OPTI_D1_RESTORE_FROM)
# opti_state_dict_d2 = torch.load(OPTI_D2_RESTORE_FROM)
# optimizer.load_state_dict(opti_state_dict)
# optimizer_D1.load_state_dict(opti_state_dict_d1)
# optimizer_D1.load_state_dict(opti_state_dict_d2)
optimizer.zero_grad()
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
interp_target = nn.Upsample(size=(input_size_target[1],
input_size_target[0]),
mode='bilinear',
align_corners=True)
# labels for adversarial training
source_label = 1
target_label = 0
mIoUs = []
i_iters = []
for i_iter in range(args.num_steps):
if i_iter <= iter_start:
continue
loss_seg_value1 = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_seg_value2 = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate_D(optimizer_D1, i_iter)
adjust_learning_rate_D(optimizer_D2, i_iter)
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D1.parameters():
param.requires_grad = False
for param in model_D2.parameters():
param.requires_grad = False
# train with source
_, batch = trainloader_iter.__next__()
images, labels, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
loss_seg1 = loss_calc(pred1, labels, args.gpu)
loss_seg2 = loss_calc(pred2, labels, args.gpu)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value1 += loss_seg1.data.cpu().numpy() / args.iter_size
loss_seg_value2 += loss_seg2.data.cpu().numpy() / args.iter_size
# train with target
_, batch = targetloader_iter.__next__()
images, _, name = batch
images = Variable(images).cuda(args.gpu)
pred_target1, pred_target2 = model(images)
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
D_out1 = model_D1(F.softmax(pred_target1, dim=1))
D_out2 = model_D2(F.softmax(pred_target2, dim=1))
loss_adv_target1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_adv_target2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda(
args.gpu))
loss = args.lambda_adv_target1 * loss_adv_target1 + args.lambda_adv_target2 * loss_adv_target2
loss = loss / args.iter_size
loss.backward()
loss_adv_target_value1 += loss_adv_target1.data.cpu().numpy(
) / args.iter_size
loss_adv_target_value2 += loss_adv_target2.data.cpu().numpy(
) / args.iter_size
# train D
# bring back requires_grad
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
# train with source
pred1 = pred1.detach()
pred2 = pred2.detach()
D_out1 = model_D1(F.softmax(pred1, dim=1))
D_out2 = model_D2(F.softmax(pred2, dim=1))
weight_s = float(D_out2.mean().data.cpu().numpy())
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
loss_D_value1 += loss_D1.data.cpu().numpy()
loss_D_value2 += loss_D2.data.cpu().numpy()
# train with target
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
D_out1 = model_D1(F.softmax(pred_target1, dim=1))
D_out2 = model_D2(F.softmax(pred_target2, dim=1))
weight_t = float(D_out2.mean().data.cpu().numpy())
# if weight_b>0.5 and i_iter>500:
# confidence_map = interp(D_out2).cpu().data[0][0].numpy()
# name = name[0].split('/')[-1]
# confidence_map=255*confidence_map
# confidence_output=Image.fromarray(confidence_map.astype(np.uint8))
# confidence_output.save('./result/confid_map/%s.png' % (name.split('.')[0]))
# zq=1
print(weight_s, weight_t)
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(target_label)).cuda(
args.gpu))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(target_label)).cuda(
args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
loss_D_value1 += loss_D1.data.cpu().numpy()
loss_D_value2 += loss_D2.data.cpu().numpy()
optimizer.step()
optimizer_D1.step()
optimizer_D2.step()
# print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f} loss_adv1 = {4:.3f}, loss_adv2 = {5:.3f} loss_D1 = {6:.3f} loss_D2 = {7:.3f}'
.format(i_iter, args.num_steps, loss_seg_value1, loss_seg_value2,
loss_adv_target_value1, loss_adv_target_value2,
loss_D_value1, loss_D_value2))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps) + '_D2.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, 'GTA5_' + str(i_iter) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D2.pth'))
torch.save(
optimizer.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(i_iter) + '_optimizer.pth'))
torch.save(
optimizer_D1.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(i_iter) + '_optimizer_D1.pth'))
torch.save(
optimizer_D2.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(i_iter) + '_optimizer_D2.pth'))
show_pred_sv_dir = pre_sv_dir.format(i_iter)
mIoU = show_val(model.state_dict(), show_pred_sv_dir, gpu)
mIoUs.append(str(round(np.nanmean(mIoU) * 100, 2)))
i_iters.append(i_iter)
print_i = 0
for miou in mIoUs:
print('i{0}: {1}'.format(i_iters[print_i], miou))
print_i = print_i + 1
def main():
"""Create the model and start the training."""
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
h, w = map(int, args.input_size_target.split(','))
input_size_target = (h, w)
cudnn.enabled = True
from pytorchgo.utils.pytorch_utils import set_gpu
set_gpu(args.gpu)
# Create network
if args.model == 'DeepLab':
logger.info("adopting Deeplabv2 base model..")
model = Res_Deeplab(num_classes=args.num_classes, multi_scale=False)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
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 == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# print i_parts
model.load_state_dict(new_params)
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.model == "FCN8S":
logger.info("adopting FCN8S base model..")
from pytorchgo.model.MyFCN8s import MyFCN8s
model = MyFCN8s(n_class=NUM_CLASSES)
vgg16 = torchfcn.models.VGG16(pretrained=True)
model.copy_params_from_vgg16(vgg16)
optimizer = optim.SGD(model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
else:
raise ValueError
model.train()
model.cuda()
cudnn.benchmark = True
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes)
model_D2 = FCDiscriminator(num_classes=args.num_classes)
model_D1.train()
model_D1.cuda()
model_D2.train()
model_D2.cuda()
if SOURCE_DATA == "GTA5":
trainloader = data.DataLoader(GTA5DataSet(
args.data_dir,
args.data_list,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
elif SOURCE_DATA == "SYNTHIA":
trainloader = data.DataLoader(SynthiaDataSet(
args.data_dir,
args.data_list,
LABEL_LIST_PATH,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
else:
raise ValueError
targetloader = data.DataLoader(cityscapesDataSet(
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False,
mirror=args.random_mirror,
mean=IMG_MEAN,
set=args.set),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer.zero_grad()
optimizer_D1 = optim.Adam(model_D1.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D1.zero_grad()
optimizer_D2 = optim.Adam(model_D2.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D2.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
interp_target = nn.Upsample(size=(input_size_target[1],
input_size_target[0]),
mode='bilinear')
# labels for adversarial training
source_label = 0
target_label = 1
best_mIoU = 0
model_summary([model, model_D1, model_D2])
optimizer_summary([optimizer, optimizer_D1, optimizer_D2])
for i_iter in tqdm(range(args.num_steps_stop),
total=args.num_steps_stop,
desc="training"):
loss_seg_value1 = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_seg_value2 = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
optimizer.zero_grad()
lr = adjust_learning_rate(optimizer, i_iter)
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
lr_D1 = adjust_learning_rate_D(optimizer_D1, i_iter)
lr_D2 = adjust_learning_rate_D(optimizer_D2, i_iter)
for sub_i in range(args.iter_size):
######################### train G
# don't accumulate grads in D
for param in model_D1.parameters():
param.requires_grad = False
for param in model_D2.parameters():
param.requires_grad = False
# train with source
_, batch = trainloader_iter.next()
images, labels, _, _ = batch
images = Variable(images).cuda()
pred2 = model(images)
pred2 = interp(pred2)
loss_seg2 = loss_calc(pred2, labels)
loss = loss_seg2
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value2 += loss_seg2.data.cpu().numpy()[0] / args.iter_size
# train with target
_, batch = targetloader_iter.next()
images, _, _, _ = batch
images = Variable(images).cuda()
pred_target2 = model(images)
pred_target2 = interp_target(pred_target2)
D_out2 = model_D2(F.softmax(pred_target2))
loss_adv_target2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda())
loss = args.lambda_adv_target2 * loss_adv_target2
loss = loss / args.iter_size
loss.backward()
loss_adv_target_value2 += loss_adv_target2.data.cpu().numpy(
)[0] / args.iter_size
################################## train D
# bring back requires_grad
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
# train with source
pred2 = pred2.detach()
D_out2 = model_D2(F.softmax(pred2))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda())
loss_D2 = loss_D2 / args.iter_size / 2
loss_D2.backward()
loss_D_value2 += loss_D2.data.cpu().numpy()[0]
# train with target
pred_target2 = pred_target2.detach()
D_out2 = model_D2(F.softmax(pred_target2))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(target_label)).cuda())
loss_D2 = loss_D2 / args.iter_size / 2
loss_D2.backward()
loss_D_value2 += loss_D2.data.cpu().numpy()[0]
optimizer.step()
optimizer_D1.step()
optimizer_D2.step()
if i_iter % 100 == 0:
logger.info(
'iter = {}/{},loss_seg1 = {:.3f} loss_seg2 = {:.3f} loss_adv1 = {:.3f}, loss_adv2 = {:.3f} loss_D1 = {:.3f} loss_D2 = {:.3f}, lr={:.7f}, lr_D={:.7f}, best miou16= {:.5f}'
.format(i_iter, args.num_steps_stop, loss_seg_value1,
loss_seg_value2, loss_adv_target_value1,
loss_adv_target_value2, loss_D_value1, loss_D_value2,
lr, lr_D1, best_mIoU))
if i_iter % args.save_pred_every == 0 and i_iter != 0:
logger.info("saving snapshot.....")
cur_miou16 = proceed_test(model, input_size)
is_best = True if best_mIoU < cur_miou16 else False
if is_best:
best_mIoU = cur_miou16
torch.save(
{
'iteration': i_iter,
'optim_state_dict': optimizer.state_dict(),
'optim_D1_state_dict': optimizer_D1.state_dict(),
'optim_D2_state_dict': optimizer_D2.state_dict(),
'model_state_dict': model.state_dict(),
'model_D1_state_dict': model_D1.state_dict(),
'model_D2_state_dict': model_D2.state_dict(),
'best_mean_iu': cur_miou16,
}, osp.join(logger.get_logger_dir(), 'checkpoint.pth.tar'))
if is_best:
import shutil
shutil.copy(
osp.join(logger.get_logger_dir(), 'checkpoint.pth.tar'),
osp.join(logger.get_logger_dir(), 'model_best.pth.tar'))
if i_iter >= args.num_steps_stop - 1:
break
def main():
"""Create the model and start the training."""
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
h, w = map(int, args.input_size_target.split(','))
input_size_target = (h, w)
cudnn.enabled = True
gpu = args.gpu
# Create network
if args.model == 'DeepLab':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
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 == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# print i_parts
model.load_state_dict(new_params)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes)
model_D2 = FCDiscriminator(num_classes=args.num_classes)
model_D1.train()
model_D1.cuda(args.gpu)
model_D2.train()
model_D2.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(SynthiaDataSet(
args.data_dir,
args.data_list,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(CamvidDataSet(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False,
mirror=args.random_mirror,
mean=IMG_MEAN,
set=args.set),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D1 = optim.Adam(model_D1.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D1.zero_grad()
optimizer_D2 = optim.Adam(model_D2.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D2.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
interp_target = nn.Upsample(size=(input_size_target[1],
input_size_target[0]),
mode='bilinear',
align_corners=True)
# labels for adversarial training
source_label = 0
target_label = 1
for i_iter in range(args.num_steps):
loss_seg_value1 = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_seg_value2 = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate_D(optimizer_D1, i_iter)
adjust_learning_rate_D(optimizer_D2, i_iter)
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D1.parameters():
param.requires_grad = False
for param in model_D2.parameters():
param.requires_grad = False
# train with source
_, batch = trainloader_iter.next()
images, labels, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
loss_seg1 = loss_calc(pred1, labels, args.gpu)
loss_seg2 = loss_calc(pred2, labels, args.gpu)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value1 += loss_seg1.data.cpu().numpy() / args.iter_size
loss_seg_value2 += loss_seg2.data.cpu().numpy() / args.iter_size
# train with target
_, batch = targetloader_iter.next()
images, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred_target1, pred_target2 = model(images)
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_adv_target1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_adv_target2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda(
args.gpu))
loss = args.lambda_adv_target1 * loss_adv_target1 + args.lambda_adv_target2 * loss_adv_target2
loss = loss / args.iter_size
loss.backward()
loss_adv_target_value1 += loss_adv_target1.data.cpu().numpy(
) / args.iter_size
loss_adv_target_value2 += loss_adv_target2.data.cpu().numpy(
) / args.iter_size
# train D
# bring back requires_grad
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
# train with source
pred1 = pred1.detach()
pred2 = pred2.detach()
D_out1 = model_D1(F.softmax(pred1))
D_out2 = model_D2(F.softmax(pred2))
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
loss_D_value1 += loss_D1.data.cpu().numpy()
loss_D_value2 += loss_D2.data.cpu().numpy()
# train with target
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(target_label)).cuda(
args.gpu))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(target_label)).cuda(
args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
loss_D_value1 += loss_D1.data.cpu().numpy()
loss_D_value2 += loss_D2.data.cpu().numpy()
optimizer.step()
optimizer_D1.step()
optimizer_D2.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f} loss_adv1 = {4:.3f}, loss_adv2 = {5:.3f} loss_D1 = {6:.3f} loss_D2 = {7:.3f}'
.format(i_iter, args.num_steps, loss_seg_value1, loss_seg_value2,
loss_adv_target_value1, loss_adv_target_value2,
loss_D_value1, loss_D_value2))
if i_iter >= args.num_steps_stop - 1:
print 'save model ...'
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps) + '_D2.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, 'GTA5_' + str(i_iter) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D2.pth'))
def main():
"""Create the model and start the training."""
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
h, w = map(int, args.com_size.split(','))
com_size = (h, w)
h, w = map(int, args.input_size_target.split(','))
input_size_target = (h, w)
cudnn.enabled = True
gpu = args.gpu
torch.cuda.set_device(args.gpu)
# Create network
if args.model == 'DeepLab':
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(
args.restore_from,
map_location=lambda storage, loc: storage.cuda(args.gpu))
model.load_state_dict(saved_state_dict)
# Create network
# if args.model == 'DeepLab':
# model = Res_Deeplab(num_classes=args.num_classes)
# # saved_state_dict = torch.load(args.restore_from, map_location=lambda storage, loc: storage.cuda(args.gpu))
# saved_state_dict = torch.loa
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
print("amy", torch.cuda.current_device())
############################
#validation data
testloader = data.DataLoader(dataset.cityscapes_dataset.cityscapesDataSet(
args.data_dir_target,
args.data_list_target_val,
crop_size=input_size_target,
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set_val),
batch_size=1,
shuffle=False,
pin_memory=True)
with open('./dataset/cityscapes_list/info.json', 'r') as fp:
info = json.load(fp)
mapping = np.array(info['label2train'], dtype=np.int)
label_path_list = './dataset/cityscapes_list/label.txt'
gt_imgs = open(label_path_list, 'r').read().splitlines()
gt_imgs = [join('./data/Cityscapes/data/gtFine/val', x) for x in gt_imgs]
interp_val = nn.Upsample(size=(com_size[1], com_size[0]), mode='bilinear')
############################
trainloader = data.DataLoader(GTA5DataSet(args.data_dir,
args.data_list,
max_iters=args.num_steps *
args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False,
mirror=args.random_mirror,
mean=IMG_MEAN,
set=args.set),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
# interp_target = nn.UpsamplingBilinear2d(size=(input_size_target[1], input_size_target[0]))
Softmax = torch.nn.Softmax()
bce_loss = torch.nn.BCEWithLogitsLoss()
for i_iter in range(args.num_steps):
# loss_seg_value1 = 0
loss_seg_value = 0
loss_weak_value = 0
loss_neg_value = 0
loss_lse_source_value = 0
loss_lse_target_value = 0
entropy_samples_value = 0
model.train()
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
for sub_i in range(args.iter_size):
# train with source
# _, batch = next(trainloader_iter)
# images, labels, class_label_source, mask_weakly, _, name = batch
# images = Variable(images).cuda(args.gpu)
# pred = model(images)
# pred = interp(pred)
# loss_seg = loss_calc(pred, labels, args.gpu)
# num = torch.sum(mask_weakly[0][0]).data.item()
# class_label_source_lse = class_label_source.type(torch.FloatTensor)
# exp_source = torch.min(torch.exp(1*pred), Variable(torch.exp(torch.tensor(40.0))).cuda(args.gpu))
# lse = (1.0/1) * torch.log( (512*256/num) * AvePool(torch.exp(1*pred) * mask_weakly.type(torch.FloatTensor).cuda(args.gpu)))
# loss_lse_source = bce_loss(lse, Variable(class_label_source_lse.reshape(lse.
_, batch = next(targetloader_iter)
images, class_label, _, _ = batch
images = Variable(images).cuda(args.gpu)
_, pred_target = model(images)
# optimizer.step()
del pred_target, batch, images
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f} loss_lse_source = {3:.3f} loss_lse_target = {4:.3f}'
.format(i_iter, args.num_steps, loss_seg_value,
loss_lse_source_value, loss_lse_target_value))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
model.eval()
hist = np.zeros((19, 19))
f = open(args.results_dir, 'a')
for index, batch in enumerate(testloader):
print(index)
image, _, name = batch
_, output = model(
Variable(image, volatile=True).cuda(args.gpu))
pred = interp_val(output)
pred = pred[0].permute(1, 2, 0)
pred = torch.max(pred, 2)[1].byte()
pred_cpu = pred.data.cpu().numpy()
del pred, output
label = Image.open(gt_imgs[index])
label = np.array(label.resize(com_size, Image.NEAREST))
label = label_mapping(label, mapping)
hist += fast_hist(label.flatten(), pred_cpu.flatten(), 19)
mIoUs = per_class_iu(hist)
mIoU = round(np.nanmean(mIoUs) * 100, 2)
print(mIoU)
f.write('i_iter:{:d}, miou:{:0.5f} \n'.format(i_iter, mIoU))
f.close()
def main():
"""Create the model and start the training."""
model_num = 0
torch.manual_seed(args.random_seed)
torch.cuda.manual_seed_all(args.random_seed)
random.seed(args.random_seed)
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
h, w = map(int, args.input_size_target.split(','))
input_size_target = (h, w)
cudnn.enabled = True
gpu = args.gpu
# Create network
if args.model == 'DeepLab':
if args.training_option == 1:
model = Res_Deeplab(num_classes=args.num_classes,
num_layers=args.num_layers)
elif args.training_option == 2:
model = Res_Deeplab2(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
for k, v in saved_state_dict.items():
print(k)
for k in new_params:
print(k)
for i in saved_state_dict:
i_parts = i.split('.')
if '.'.join(i_parts[args.i_parts_index:]) in new_params:
print("Restored...")
if args.not_restore_last == True:
if not i_parts[
args.i_parts_index] == 'layer5' and not i_parts[
args.i_parts_index] == 'layer6':
new_params['.'.join(i_parts[args.i_parts_index:]
)] = saved_state_dict[i]
else:
new_params['.'.join(
i_parts[args.i_parts_index:])] = saved_state_dict[i]
model.load_state_dict(new_params)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
writer = SummaryWriter(log_dir=args.snapshot_dir)
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes)
model_D2 = FCDiscriminator(num_classes=args.num_classes)
model_D1.train()
model_D1.cuda(args.gpu)
model_D2.train()
model_D2.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
'''trainloader = data.DataLoader(sourceDataSet(args.data_dir,
args.data_list,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN_SOURCE,
ignore_label=args.ignore_label),
batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True)'''
trainloader = data.DataLoader(sourceDataSet(
args.data_dir,
args.data_list,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
random_rotate=args.augment_1,
random_flip=args.augment_1,
random_lighting=args.augment_1,
mean=IMG_MEAN_SOURCE,
ignore_label=args.ignore_label),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(isprsDataSet(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False,
mirror=args.random_mirror,
mean=IMG_MEAN_TARGET,
ignore_label=args.ignore_label),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
valloader = data.DataLoader(valDataSet(args.data_dir_val,
args.data_list_val,
crop_size=input_size_target,
mean=IMG_MEAN_TARGET,
scale=False,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D1 = optim.Adam(model_D1.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D1.zero_grad()
optimizer_D2 = optim.Adam(model_D2.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D2.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
# EDITTED by me
interp = nn.Upsample(size=(input_size[0], input_size[1]), mode='bilinear')
interp_target = nn.Upsample(size=(input_size_target[0],
input_size_target[1]),
mode='bilinear')
# labels for adversarial training
source_label = 0
target_label = 1
# Which layers to freeze
non_trainable(args.dont_train, model)
for i_iter in range(args.num_steps):
loss_seg_value1 = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_seg_value2 = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate_D(optimizer_D1, i_iter)
adjust_learning_rate_D(optimizer_D2, i_iter)
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D1.parameters():
param.requires_grad = False
for param in model_D2.parameters():
param.requires_grad = False
# train with source
while True:
try:
_, batch = next(trainloader_iter)
images, labels, _, train_name = batch
#print(train_name)
images = Variable(images).cuda(args.gpu)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
# Save img
'''if i_iter % 5 == 0:
save_image_for_test(concatenate_side_by_side([images, labels, pred2]), i_iter)'''
loss_seg1 = loss_calc(pred1, labels, args.gpu,
args.ignore_label, train_name)
loss_seg2 = loss_calc(pred2, labels, args.gpu,
args.ignore_label, train_name)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# proper normalization
loss = loss / args.iter_size
loss.backward()
if isinstance(loss_seg1.data.cpu().numpy(), list):
loss_seg_value1 += loss_seg1.data.cpu().numpy(
)[0] / args.iter_size
else:
loss_seg_value1 += loss_seg1.data.cpu().numpy(
) / args.iter_size
if isinstance(loss_seg2.data.cpu().numpy(), list):
loss_seg_value2 += loss_seg2.data.cpu().numpy(
)[0] / args.iter_size
else:
loss_seg_value2 += loss_seg2.data.cpu().numpy(
) / args.iter_size
break
except (RuntimeError, AssertionError, AttributeError):
continue
if args.experiment == 1:
# Which layers to freeze
non_trainable('0', model)
# train with target
_, batch = next(targetloader_iter)
images, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred_target1, pred_target2 = model(images)
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
#total_image2 = vutils.make_grid(torch.cat((images.cuda()), dim = 2),normalize=True, scale_each=True)
#total_image2 = images.cuda()
#, pred_target1.cuda(), pred_target2.cuda()
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_adv_target1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_adv_target2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda(
args.gpu))
loss = args.lambda_adv_target1 * loss_adv_target1 + args.lambda_adv_target2 * loss_adv_target2
loss = loss / args.iter_size
loss.backward()
if isinstance(loss_adv_target1.data.cpu().numpy(), list):
loss_adv_target_value1 += loss_adv_target1.data.cpu().numpy(
)[0] / args.iter_size
else:
loss_adv_target_value1 += loss_adv_target1.data.cpu().numpy(
) / args.iter_size
if isinstance(loss_adv_target2.data.cpu().numpy(), list):
loss_adv_target_value2 += loss_adv_target2.data.cpu().numpy(
)[0] / args.iter_size
else:
loss_adv_target_value2 += loss_adv_target2.data.cpu().numpy(
) / args.iter_size
if args.experiment == 1:
# Which layers to freeze
non_trainable(args.dont_train, model)
# train D
# bring back requires_grad
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
# train with source
pred1 = pred1.detach()
pred2 = pred2.detach()
D_out1 = model_D1(F.softmax(pred1))
D_out2 = model_D2(F.softmax(pred2))
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
if isinstance(loss_D1.data.cpu().numpy(), list):
loss_D_value1 += loss_D1.data.cpu().numpy()[0]
else:
loss_D_value1 += loss_D1.data.cpu().numpy()
if isinstance(loss_D2.data.cpu().numpy(), list):
loss_D_value2 += loss_D2.data.cpu().numpy()[0]
else:
loss_D_value2 += loss_D2.data.cpu().numpy()
# train with target
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(target_label)).cuda(
args.gpu))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(target_label)).cuda(
args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
if isinstance(loss_D1.data.cpu().numpy(), list):
loss_D_value1 += loss_D1.data.cpu().numpy()[0]
else:
loss_D_value1 += loss_D1.data.cpu().numpy()
if isinstance(loss_D2.data.cpu().numpy(), list):
loss_D_value2 += loss_D2.data.cpu().numpy()[0]
else:
loss_D_value2 += loss_D2.data.cpu().numpy()
optimizer.step()
optimizer_D1.step()
optimizer_D2.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f} loss_adv1 = {4:.3f}, loss_adv2 = {5:.3f} loss_D1 = {6:.3f} loss_D2 = {7:.3f}'
.format(i_iter, args.num_steps, loss_seg_value1, loss_seg_value2,
loss_adv_target_value1, loss_adv_target_value2,
loss_D_value1, loss_D_value2))
if i_iter >= args.num_steps_stop - 1:
#print ('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'model_' + str(args.num_steps) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'model_' + str(args.num_steps) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir,
'model_' + str(args.num_steps) + '_D2.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, 'model_' + str(i_iter) + '.pth'))
torch.save(model_D1.state_dict(), osp.join(args.snapshot_dir, 'model_' + str(i_iter) + '_D1.pth'))
torch.save(model_D2.state_dict(), osp.join(args.snapshot_dir, 'model_' + str(i_iter) + '_D2.pth'))'''
if model_num != args.num_models_keep:
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'model_' + str(model_num) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'model_' + str(model_num) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir,
'model_' + str(model_num) + '_D2.pth'))
model_num = model_num + 1
if model_num == args.num_models_keep:
model_num = 0
# Validation
if (i_iter % args.val_every == 0 and i_iter != 0) or i_iter == 1:
validation(valloader, model, interp_target, writer, i_iter,
[37, 41, 10])
# Save for tensorboardx
writer.add_scalar('loss_seg_value1', loss_seg_value1, i_iter)
writer.add_scalar('loss_seg_value2', loss_seg_value2, i_iter)
writer.add_scalar('loss_adv_target_value1', loss_adv_target_value1,
i_iter)
writer.add_scalar('loss_adv_target_value2', loss_adv_target_value2,
i_iter)
writer.add_scalar('loss_D_value1', loss_D_value1, i_iter)
writer.add_scalar('loss_D_value2', loss_D_value2, i_iter)
writer.close()
def main():
"""Create the model and start the training."""
w, h = map(int, args.input_size.split(','))
input_size = (w, h)
w, h = map(int, args.input_size_target.split(','))
input_size_target = (w, h)
h, w = map(int, args.com_size.split(','))
com_size = (h, w)
############################
#validation data
testloader = data.DataLoader(cityscapesDataSet(args.data_dir_target, args.data_list_target_val, crop_size=input_size, mean=IMG_MEAN, scale=False, mirror=False, set=args.set_val),
batch_size=1, shuffle=False, pin_memory=True)
with open('./dataset/cityscapes_list/info.json', 'r') as fp:
info = json.load(fp)
mapping = np.array(info['label2train'], dtype=np.int)
label_path_list = './dataset/cityscapes_list/label.txt'
gt_imgs = open(label_path_list, 'r').read().splitlines()
gt_imgs = [osp.join('./data/Cityscapes/data/gtFine/val', x) for x in gt_imgs]
interp_val = nn.UpsamplingBilinear2d(size=(com_size[1], com_size[0]))
############################
cudnn.enabled = True
# Create network
# if args.model == 'DeepLab':
# model = Res_Deeplab(num_classes=args.num_classes)
# if args.restore_from[:4] == 'http' :
# saved_state_dict = model_zoo.load_url(args.restore_from)
# else:
# 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 == 19 or not i_parts[1] == 'layer5':
# new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# # print i_parts
# model.load_state_dict(new_params)
if args.model == 'DeepLab':
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes)
model_D2 = FCDiscriminator(num_classes=args.num_classes)
model_D1.train()
model_D1.cuda(args.gpu)
model_D2.train()
model_D2.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(
GTA5DataSet(args.data_dir, args.data_list, max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN),
batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(cityscapesDataSet(args.data_dir_target, args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False, mirror=args.random_mirror, mean=IMG_MEAN,
set=args.set),
batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D1 = optim.Adam(model_D1.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
optimizer_D1.zero_grad()
optimizer_D2 = optim.Adam(model_D2.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
optimizer_D2.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear')
# labels for adversarial training
source_label = 0
target_label = 1
AvePool = torch.nn.AvgPool2d(kernel_size=(512,1024))
for i_iter in range(args.num_steps):
model.train()
loss_lse_target_value = 0
loss_seg_value1 = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_seg_value2 = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate_D(optimizer_D1, i_iter)
adjust_learning_rate_D(optimizer_D2, i_iter)
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D1.parameters():
param.requires_grad = False
for param in model_D2.parameters():
param.requires_grad = False
# train with source
_, batch = next(trainloader_iter)
images, labels, class_label_source, mask_weakly, _, name = batch
images = Variable(images).cuda(args.gpu)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
loss_seg1 = loss_calc(pred1, labels, args.gpu)
loss_seg2 = loss_calc(pred2, labels, args.gpu)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value1 += loss_seg1.data.item() / args.iter_size
loss_seg_value2 += loss_seg2.data.item() / args.iter_size
# train with target
_, batch = next(targetloader_iter)
images, class_label, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred_target1, pred_target2 = model(images)
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
class_label_target_lse = class_label.type(torch.FloatTensor)
exp_target = torch.min(torch.exp(1*pred_target2), Variable(torch.exp(torch.tensor(40.0))).cuda(args.gpu))
lse = (1.0/1) * torch.log(AvePool(exp_target))
loss_lse_target = bce_loss(lse, Variable(class_label_target_lse.reshape(lse.size())).cuda(args.gpu))
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_adv_target1 = bce_loss(D_out1,
Variable(torch.FloatTensor(D_out1.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_adv_target2 = bce_loss(D_out2,
Variable(torch.FloatTensor(D_out2.data.size()).fill_(source_label)).cuda(
args.gpu))
loss = args.lambda_adv_target1 * loss_adv_target1 + args.lambda_adv_target2 * loss_adv_target2 + 0.2 * loss_lse_target
loss = loss / args.iter_size
loss.backward()
loss_adv_target_value1 += loss_adv_target1.data.item() / args.iter_size
loss_adv_target_value2 += loss_adv_target2.data.item() / args.iter_size
loss_lse_target_value += loss_lse_target.data.item() / args.iter_size
# train D
# bring back requires_grad
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
# train with source
pred1 = pred1.detach()
pred2 = pred2.detach()
D_out1 = model_D1(F.softmax(pred1))
D_out2 = model_D2(F.softmax(pred2))
loss_D1 = bce_loss(D_out1,
Variable(torch.FloatTensor(D_out1.data.size()).fill_(source_label)).cuda(args.gpu))
loss_D2 = bce_loss(D_out2,
Variable(torch.FloatTensor(D_out2.data.size()).fill_(source_label)).cuda(args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
loss_D_value1 += loss_D1.data.item()
loss_D_value2 += loss_D2.data.item()
# train with target
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_D1 = bce_loss(D_out1,
Variable(torch.FloatTensor(D_out1.data.size()).fill_(target_label)).cuda(args.gpu))
loss_D2 = bce_loss(D_out2,
Variable(torch.FloatTensor(D_out2.data.size()).fill_(target_label)).cuda(args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
loss_D_value1 += loss_D1.data.item()
loss_D_value2 += loss_D2.data.item()
optimizer.step()
optimizer_D1.step()
optimizer_D2.step()
del D_out1, D_out2, pred1, pred2, pred_target1, pred_target2, images, labels
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f} loss_adv1 = {4:.3f}, loss_adv2 = {5:.3f} loss_D1 = {6:.3f} loss_D2 = {7:.3f} loss_lse_target = {8:.3f}'.format(i_iter, args.num_steps, loss_seg_value1, loss_seg_value2, loss_adv_target_value1, loss_adv_target_value2, loss_D_value1, loss_D_value2, loss_lse_target_value))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'GTA5_' + str(args.num_steps_stop) + '.pth'))
torch.save(model_D1.state_dict(), osp.join(args.snapshot_dir, 'GTA5_' + str(args.num_steps_stop) + '_D1.pth'))
torch.save(model_D2.state_dict(), osp.join(args.snapshot_dir, 'GTA5_' + str(args.num_steps_stop) + '_D2.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, 'GTA5_' + str(i_iter) + '.pth'))
hist = np.zeros((19, 19))
# model.cuda(0)
model.eval()
f = open(args.results_dir, 'a')
for index, batch in enumerate(testloader):
print(index)
image, _, _, name = batch
output1, output2 = model(Variable(image, volatile=True).cuda(args.gpu))
pred = interp_val(output2)
pred = pred[0].permute(1,2,0)
pred = torch.max(pred, 2)[1].byte()
pred_cpu = pred.data.cpu().numpy()
del pred, output1, output2
label = Image.open(gt_imgs[index])
label = np.array(label.resize(com_size, Image.NEAREST))
label = label_mapping(label, mapping)
hist += fast_hist(label.flatten(), pred_cpu.flatten(), 19)
# model.cuda(args.gpu)
mIoUs = per_class_iu(hist)
mIoU = round(np.nanmean(mIoUs) * 100, 2)
print(mIoU)
f.write('i_iter:{:d}, miou:{:0.5f} \n'.format(i_iter,mIoU))
f.close()
def main():
"""Create the model and start the training."""
model_num = 0 # The number of model (for saving models)
torch.manual_seed(args.random_seed)
torch.cuda.manual_seed_all(args.random_seed)
random.seed(args.random_seed)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
writer = SummaryWriter(log_dir=args.snapshot_dir)
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
h, w = map(int, args.input_size_target.split(','))
input_size_target = (h, w)
cudnn.enabled = True
gpu = args.gpu
cudnn.benchmark = True
# init G
if args.model == 'DeepLab':
if args.training_option == 1:
model = Res_Deeplab(num_classes=args.num_classes,
num_layers=args.num_layers,
dropout=args.dropout,
after_layer=args.after_layer)
elif args.training_option == 2:
model = Res_Deeplab2(num_classes=args.num_classes)
'''elif args.training_option == 3:
model = Res_Deeplab50(num_classes=args.num_classes)'''
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
for k, v in saved_state_dict.items():
print(k)
for k in new_params:
print(k)
for i in saved_state_dict:
i_parts = i.split('.')
if '.'.join(i_parts[args.i_parts_index:]) in new_params:
print("Restored...")
if args.not_restore_last == True:
if not i_parts[
args.i_parts_index] == 'layer5' and not i_parts[
args.i_parts_index] == 'layer6':
new_params['.'.join(i_parts[args.i_parts_index:]
)] = saved_state_dict[i]
else:
new_params['.'.join(
i_parts[args.i_parts_index:])] = saved_state_dict[i]
model.load_state_dict(new_params)
model.train()
model.cuda(args.gpu)
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes,
extra_layers=args.extra_discriminator_layers)
model_D2 = FCDiscriminator(num_classes=args.num_classes, extra_layers=0)
model_D1.train()
model_D1.cuda(args.gpu)
model_D2.train()
model_D2.cuda(args.gpu)
trainloader = data.DataLoader(sourceDataSet(
args.data_dir,
args.data_list,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
random_rotate=False,
random_flip=args.augment_1,
random_lighting=args.augment_1,
random_blur=args.augment_1,
random_scaling=args.augment_1,
mean=IMG_MEAN_SOURCE,
ignore_label=args.ignore_label),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
trainloader2 = data.DataLoader(sourceDataSet(
args.data_dir2,
args.data_list2,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
random_rotate=False,
random_flip=args.augment_2,
random_lighting=args.augment_2,
random_blur=args.augment_2,
random_scaling=args.augment_2,
mean=IMG_MEAN_SOURCE2,
ignore_label=args.ignore_label),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter2 = enumerate(trainloader2)
if args.num_of_targets > 1:
IMG_MEAN_TARGET1 = np.array(
(101.41694189393208, 89.68194541655483, 77.79408426901315),
dtype=np.float32) # crowdai all BGR
targetloader1 = data.DataLoader(isprsDataSet(
args.data_dir_target1,
args.data_list_target1,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False,
mean=IMG_MEAN_TARGET1,
ignore_label=args.ignore_label),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter1 = enumerate(targetloader1)
targetloader = data.DataLoader(isprsDataSet(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False,
mean=IMG_MEAN_TARGET,
ignore_label=args.ignore_label),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
valloader = data.DataLoader(valDataSet(args.data_dir_val,
args.data_list_val,
crop_size=input_size_target,
mean=IMG_MEAN_TARGET,
scale=args.val_scale,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D1 = optim.Adam(model_D1.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D1.zero_grad()
optimizer_D2 = optim.Adam(model_D2.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D2.zero_grad()
if args.weighted_loss == True:
bce_loss = torch.nn.BCEWithLogitsLoss()
else:
bce_loss = torch.nn.BCEWithLogitsLoss()
interp = nn.Upsample(size=(input_size[0], input_size[1]), mode='bilinear')
interp_target = nn.Upsample(size=(input_size_target[0],
input_size_target[1]),
mode='bilinear')
# labels for adversarial training
source_label = 0
target_label = 1
# Which layers to freeze
non_trainable(args.dont_train, model)
# List saving all best 5 mIoU's
best_mIoUs = [0.0, 0.0, 0.0, 0.0, 0.0]
for i_iter in range(args.num_steps):
loss_seg_value1 = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_seg_value2 = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
optimizer.zero_grad()
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate(optimizer, i_iter)
adjust_learning_rate_D(optimizer_D1, i_iter)
adjust_learning_rate_D(optimizer_D2, i_iter)
for sub_i in range(args.iter_size):
################################## train G #################################
# don't accumulate grads in D
for param in model_D1.parameters():
param.requires_grad = False
for param in model_D2.parameters():
param.requires_grad = False
################################## train with source #################################
while True:
try:
_, batch = next(
trainloader_iter) # Cityscapes, only discriminator1
images, labels, _, train_name = batch
images = Variable(images).cuda(args.gpu)
_, batch = next(
trainloader_iter2
) # Main (airsim) discriminator2 and final output
images2, labels2, size, train_name2 = batch
images2 = Variable(images2).cuda(args.gpu)
pred1, _ = model(images)
pred1 = interp(pred1)
_, pred2 = model(images2)
pred2 = interp(pred2)
loss_seg1 = loss_calc(pred1, labels, args.gpu,
args.ignore_label, train_name,
weights1)
loss_seg2 = loss_calc(pred2, labels2, args.gpu,
args.ignore_label, train_name2,
weights2)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# proper normalization
loss = loss / args.iter_size
loss.backward()
if isinstance(loss_seg1.data.cpu().numpy(), list):
loss_seg_value1 += loss_seg1.data.cpu().numpy(
)[0] / args.iter_size
else:
loss_seg_value1 += loss_seg1.data.cpu().numpy(
) / args.iter_size
if isinstance(loss_seg2.data.cpu().numpy(), list):
loss_seg_value2 += loss_seg2.data.cpu().numpy(
)[0] / args.iter_size
else:
loss_seg_value2 += loss_seg2.data.cpu().numpy(
) / args.iter_size
break
except (RuntimeError, AssertionError, AttributeError):
continue
###################################################################################################
_, batch = next(targetloader_iter)
images, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred_target1, pred_target2 = model(images)
if args.num_of_targets > 1:
_, batch1 = next(targetloader_iter1)
images1, _, _ = batch1
images1 = Variable(images1).cuda(args.gpu)
pred_target1, _ = model(images1)
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
################################## train with target #################################
if args.adv_option == 1 or args.adv_option == 3:
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_adv_target1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_adv_target2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda(
args.gpu))
loss = args.lambda_adv_target1 * loss_adv_target1 + args.lambda_adv_target2 * loss_adv_target2
loss = loss / args.iter_size
loss.backward()
if isinstance(loss_adv_target1.data.cpu().numpy(), list):
loss_adv_target_value1 += loss_adv_target1.data.cpu(
).numpy()[0] / args.iter_size
else:
loss_adv_target_value1 += loss_adv_target1.data.cpu(
).numpy() / args.iter_size
if isinstance(loss_adv_target2.data.cpu().numpy(), list):
loss_adv_target_value2 += loss_adv_target2.data.cpu(
).numpy()[0] / args.iter_size
else:
loss_adv_target_value2 += loss_adv_target2.data.cpu(
).numpy() / args.iter_size
###################################################################################################
if args.adv_option == 2 or args.adv_option == 3:
pred1, _ = model(images)
pred1 = interp(pred1)
_, pred2 = model(images2)
pred2 = interp(pred2)
'''pred1 = pred1.detach()
pred2 = pred2.detach()'''
D_out1 = model_D1(F.softmax(pred1, dim=1))
D_out2 = model_D2(F.softmax(pred2, dim=1))
loss_adv_target1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(target_label)).cuda(
args.gpu))
loss_adv_target2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(target_label)).cuda(
args.gpu))
loss = args.lambda_adv_target1 * loss_adv_target1 + args.lambda_adv_target2 * loss_adv_target2
loss = loss / args.iter_size
loss.backward()
if isinstance(loss_adv_target1.data.cpu().numpy(), list):
loss_adv_target_value1 += loss_adv_target1.data.cpu(
).numpy()[0] / args.iter_size
else:
loss_adv_target_value1 += loss_adv_target1.data.cpu(
).numpy() / args.iter_size
if isinstance(loss_adv_target2.data.cpu().numpy(), list):
loss_adv_target_value2 += loss_adv_target2.data.cpu(
).numpy()[0] / args.iter_size
else:
loss_adv_target_value2 += loss_adv_target2.data.cpu(
).numpy() / args.iter_size
###################################################################################################
################################## train D #################################
# bring back requires_grad
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
################################## train with source #################################
pred1 = pred1.detach()
pred2 = pred2.detach()
D_out1 = model_D1(F.softmax(pred1))
D_out2 = model_D2(F.softmax(pred2))
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
if isinstance(loss_D1.data.cpu().numpy(), list):
loss_D_value1 += loss_D1.data.cpu().numpy()[0]
else:
loss_D_value1 += loss_D1.data.cpu().numpy()
if isinstance(loss_D2.data.cpu().numpy(), list):
loss_D_value2 += loss_D2.data.cpu().numpy()[0]
else:
loss_D_value2 += loss_D2.data.cpu().numpy()
################################# train with target #################################
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(target_label)).cuda(
args.gpu))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(target_label)).cuda(
args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
if isinstance(loss_D1.data.cpu().numpy(), list):
loss_D_value1 += loss_D1.data.cpu().numpy()[0]
else:
loss_D_value1 += loss_D1.data.cpu().numpy()
if isinstance(loss_D2.data.cpu().numpy(), list):
loss_D_value2 += loss_D2.data.cpu().numpy()[0]
else:
loss_D_value2 += loss_D2.data.cpu().numpy()
optimizer.step()
optimizer_D1.step()
optimizer_D2.step()
if i_iter % args.save_pred_every == 0 and i_iter != 0:
if model_num != args.num_models_keep:
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'model_' + str(model_num) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'model_' + str(model_num) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir,
'model_' + str(model_num) + '_D2.pth'))
model_num = model_num + 1
if model_num == args.num_models_keep:
model_num = 0
# Validation
if (i_iter % args.val_every == 0 and i_iter != 0) or i_iter == 1:
mIoU = validation(valloader, model, interp_target, writer, i_iter,
[37, 41, 10])
for i in range(0, len(best_mIoUs)):
if best_mIoUs[i] < mIoU:
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'bestmodel_' + str(i) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'bestmodel_' + str(i) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir,
'bestmodel_' + str(i) + '_D2.pth'))
best_mIoUs.append(mIoU)
print("Saved model at iteration %d as the best %d" %
(i_iter, i))
best_mIoUs.sort(reverse=True)
best_mIoUs = best_mIoUs[:5]
break
# Save for tensorboardx
writer.add_scalar('loss_seg_value1', loss_seg_value1, i_iter)
writer.add_scalar('loss_seg_value2', loss_seg_value2, i_iter)
writer.add_scalar('loss_adv_target_value1', loss_adv_target_value1,
i_iter)
writer.add_scalar('loss_adv_target_value2', loss_adv_target_value2,
i_iter)
writer.add_scalar('loss_D_value1', loss_D_value1, i_iter)
writer.add_scalar('loss_D_value2', loss_D_value2, i_iter)
writer.close()
def main():
"""Create the model and start the training."""
writer = SummaryWriter('./logs')
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
h, w = map(int, args.input_size_target.split(','))
input_size_target = (h, w)
cudnn.enabled = True
gpu = args.gpu
# Create network
if args.model == 'DeepLab':
model = Res_Deeplab(num_classes=args.num_classes)
model_state_dict = model.state_dict()
if args.restore_from[:4] == 'http' :
saved_state_dict = model_zoo.load_url(args.restore_from)
elif args.restore_from[:4] == 'https' :
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
saved_state_dict = {k.replace('Scale.', ''): v for k, v in saved_state_dict.items() if k.replace('Scale.', '')
in model_state_dict and 'layer5' not in k }
# 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 == 19 or not i_parts[1] == 'layer5':
# print('.'.join(i_parts[1:]),saved_state_dict[i])
# # print i_parts
# print("Key new")
# print(new_params.keys())
# print("your model new")
# print(saved_state_dict.keys())
model_state_dict.update(saved_state_dict)
model.load_state_dict(model_state_dict)
# model = torch.nn.DataParallel(model)
# patch_replication_callback(model)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# # init D
# model_D1 = FCDiscriminator(num_classes=args.num_classes)
# model_D2 = FCDiscriminator(num_classes=args.num_classes)
#
# # model_D1 = torch.nn.DataParallel(model_D1)
# # patch_replication_callback(model_D1)
# model_D1.train()
# model_D1.cuda(args.gpu)
#
# # model_D2 = torch.nn.DataParallel(model_D2)
# # patch_replication_callback(model_D2)
# model_D2.train()
# model_D2.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(
SynthiaDataSet(args.data_dir, args.data_list, max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN),
batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(cityscapesDataSet(args.data_dir_target, args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False, mirror=args.random_mirror, mean=IMG_MEAN,
set=args.set),
batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
targetloader_val = data.DataLoader(
cityscapesDataSet(args.data_dir_target, args.data_list_target.replace('train.txt', 'val.txt'),
crop_size=(1024, 512), mean=IMG_MEAN, scale=False, mirror=False, set='val'),
batch_size=1, shuffle=False,pin_memory=True)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer_D1 = optim.Adam(model_D1.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
# optimizer_D1.zero_grad()
#
# optimizer_D2 = optim.Adam(model_D2.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
# optimizer_D2.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
interp_target = nn.Upsample(size=(1024, 2048), mode='bilinear')
# labels for adversarial training
source_label = 0
target_label = 1
for i_iter in range(args.num_steps):
ipdb.set_trace()
if i_iter%(len(trainloader_iter)*args.batch_size)==0:
validate(targetloader_val, model, interp_target, writer, i_iter, args)
loss_seg_value1 = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_seg_value2 = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
# optimizer_D1.zero_grad()
# optimizer_D2.zero_grad()
# adjust_learning_rate_D(optimizer_D1, i_iter)
# adjust_learning_rate_D(optimizer_D2, i_iter)
for sub_i in range(args.iter_size):
# train G
# # don't accumulate grads in D
# for param in model_D1.parameters():
# param.requires_grad = False
#
# for param in model_D2.parameters():
# param.requires_grad = False
# train with source
_, batch = next(trainloader_iter)
images, labels, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
loss_seg1 = loss_calc(pred1, labels, args.gpu)
loss_seg2 = loss_calc(pred2, labels, args.gpu)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value1 += loss_seg1.item() / args.iter_size
loss_seg_value2 += loss_seg2.item() / args.iter_size
# # train with target
#
# _, batch = next(targetloader_iter)
# images, _, _ = batch
# images = Variable(images).cuda(args.gpu)
#
# pred_target1, pred_target2 = model(images)
# pred_target1 = interp_target(pred_target1)
# pred_target2 = interp_target(pred_target2)
#
# D_out1 = model_D1(F.softmax(pred_target1))
# D_out2 = model_D2(F.softmax(pred_target2))
#
# loss_adv_target1 = bce_loss(D_out1,
# Variable(torch.FloatTensor(D_out1.data.size()).fill_(source_label)).cuda(
# args.gpu))
#
# loss_adv_target2 = bce_loss(D_out2,
# Variable(torch.FloatTensor(D_out2.data.size()).fill_(source_label)).cuda(
# args.gpu))
#
# loss = args.lambda_adv_target1 * loss_adv_target1 + args.lambda_adv_target2 * loss_adv_target2
# loss = loss / args.iter_size
# loss.backward()
# loss_adv_target_value1 += loss_adv_target1.data.cpu().numpy()[0] / args.iter_size
# loss_adv_target_value2 += loss_adv_target2.data.cpu().numpy()[0] / args.iter_size
# # train D
#
# # bring back requires_grad
# for param in model_D1.parameters():
# param.requires_grad = True
#
# for param in model_D2.parameters():
# param.requires_grad = True
#
# # train with source
# pred1 = pred1.detach()
# pred2 = pred2.detach()
#
# D_out1 = model_D1(F.softmax(pred1))
# D_out2 = model_D2(F.softmax(pred2))
#
# loss_D1 = bce_loss(D_out1,
# Variable(torch.FloatTensor(D_out1.data.size()).fill_(source_label)).cuda(args.gpu))
#
# loss_D2 = bce_loss(D_out2,
# Variable(torch.FloatTensor(D_out2.data.size()).fill_(source_label)).cuda(args.gpu))
#
# loss_D1 = loss_D1 / args.iter_size / 2
# loss_D2 = loss_D2 / args.iter_size / 2
#
# loss_D1.backward()
# loss_D2.backward()
#
# loss_D_value1 += loss_D1.data.cpu().numpy()[0]
# loss_D_value2 += loss_D2.data.cpu().numpy()[0]
#
# # train with target
# pred_target1 = pred_target1.detach()
# pred_target2 = pred_target2.detach()
#
# D_out1 = model_D1(F.softmax(pred_target1))
# D_out2 = model_D2(F.softmax(pred_target2))
#
# loss_D1 = bce_loss(D_out1,
# Variable(torch.FloatTensor(D_out1.data.size()).fill_(target_label)).cuda(args.gpu))
#
# loss_D2 = bce_loss(D_out2,
# Variable(torch.FloatTensor(D_out2.data.size()).fill_(target_label)).cuda(args.gpu))
#
# loss_D1 = loss_D1 / args.iter_size / 2
# loss_D2 = loss_D2 / args.iter_size / 2
#
# loss_D1.backward()
# loss_D2.backward()
#
# loss_D_value1 += loss_D1.data.cpu().numpy()[0]
# loss_D_value2 += loss_D2.data.cpu().numpy()[0]
optimizer.step()
# optimizer_D1.step()
# optimizer_D2.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f} loss_adv1 = {4:.3f}, loss_adv2 = {5:.3f} loss_D1 = {6:.3f} loss_D2 = {7:.3f}'.format(
i_iter, args.num_steps, loss_seg_value1, loss_seg_value2, loss_adv_target_value1, loss_adv_target_value2, loss_D_value1, loss_D_value2))
if i_iter >= args.num_steps_stop - 1:
print ('save model ...')
torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'Synthia_' + str(args.num_steps) + '.pth'))
torch.save(model_D1.state_dict(), osp.join(args.snapshot_dir, 'Synthia_' + str(args.num_steps) + '_D1.pth'))
torch.save(model_D2.state_dict(), osp.join(args.snapshot_dir, 'Synthia_' + str(args.num_steps) + '_D2.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, 'Synthia_' + str(i_iter) + '.pth'))
torch.save(model_D1.state_dict(), osp.join(args.snapshot_dir, 'Synthia_' + str(i_iter) + '_D1.pth'))
torch.save(model_D2.state_dict(), osp.join(args.snapshot_dir, 'Synthia_' + str(i_iter) + '_D2.pth'))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
h, w = map(int, args.crop_size.split(','))
crop_size = (h, w)
h, w = map(int, args.image_size.split(','))
image_size = (h, w)
gpu0 = args.gpu
if not os.path.exists(args.save):
os.makedirs(args.save)
#model = Res_Deeplab(num_classes=args.num_classes, num_layers = args.num_layers, dropout = False)
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from[:4] == 'http' :
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.train()
model.cuda(gpu0)
testloader = data.DataLoader(valDataSet(args.data_dir,
args.data_list,
max_iters=args.num_steps,
crop_size=crop_size,
mean=IMG_MEAN,
scale=1,
mirror=False),
batch_size=1, shuffle=False, pin_memory=True)
interp = nn.Upsample(size=image_size, mode='bilinear')
data_list = []
mIoU = 0
for index, batch in enumerate(testloader):
if index % 10 == 0:
print('%d processed' % index)
image, label, size, name = batch
size = size[0].numpy()
gt = np.asarray(label[0].numpy()[:size[0],:size[1]], dtype=np.int)
_, output2 = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output2).cpu().data[0].numpy()
output = output[:,:size[0],:size[1]]
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
data_list.append([gt.flatten(), output.flatten()])
'''output_col = colorize(output, args.num_classes, palette, args.ignore_label)
label_col = colorize(np.squeeze(label, axis = 2), args.num_classes, palette, args.ignore_label)
image = image.cpu().numpy()
image = image[:,::-1,:,:]
image = np.squeeze(image, axis = 0)
image = image.transpose((1, 2, 0))
image += IMG_MEAN
image = image.transpose((2, 0, 1))
name = name[0].split('/')[-1]
#to_save = concatenate_side_by_side([image, label_col, output_col])
#to_save.save('%s/%s_eval.png' % (args.save, name.split('.')[0]))'''
mIoU = get_iou(data_list, args.num_classes)
print("Final mIoU %f" % (mIoU))
