def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(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)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testset = cityscapesDataSet(args.data_dir, args.data_list, crop_size=(1024, 512), mean=IMG_MEAN, scale=False, mirror=False, set=args.set)
testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, crop_size=(1024, 512), mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
batch_size=1, shuffle=False, pin_memory=True)
interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
for index, batch in enumerate(testloader):
if index % 100 == 0:
print '%d processd' % index
image, _, name = batch
if args.model == 'DeeplabMulti':
output1, output2 = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output2).cpu().data[0].numpy()
elif args.model == 'DeeplabVGG':
output = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
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]))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
model_path = '/path/to/checkpoint'
save_path = '/path/to/save/predictions'
#args = get_arguments()
if not os.path.exists(model_path):
print('no checkpoint')
print(model_path)
continue
if not os.path.exists(save_path):
os.makedirs(save_path)
model = Mobile_Deeplab()
saved_state_dict = torch.load(model_path)
new_params = model.state_dict().copy()
for i in saved_state_dict:
if i in new_params:
new_params[i] = saved_state_dict[i]
model.load_state_dict(new_params)
device = torch.device("cuda" if not args.cpu else "cpu")
model = model.to(device)
model.eval()
testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
batch_size=1, shuffle=False, pin_memory=True)
def test_model(model, device, val_dir, val_list, save_dir, input_size=(1024, 512)):
model.eval()
testloader = data.DataLoader(
cityscapesDataSet(val_dir, val_list, crop_size=input_size, mean=IMG_MEAN, scale=False, mirror=False,
set='val'),
batch_size=1, shuffle=False, pin_memory=True)
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True)
for index, batch in enumerate(testloader):
if index > 10:
break
image, _, name = batch
image = image.to(device)
output1, output2,_,_ = model(image)
output = interp(output2).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (save_dir, name))
output_col.save('%s/%s_color.png' % (save_dir, name.split('.')[0]))
def main():
"""Create the model and start the evaluation process."""
for i in range(1, 50):
model_path = './snapshots/GTA2Cityscapes/GTA5_{0:d}.pth'.format(i *
2000)
save_path = './result/GTA2Cityscapes_{0:d}'.format(i * 2000)
args = get_arguments()
gpu0 = args.gpu
if not os.path.exists(save_path):
os.makedirs(save_path)
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(model_path)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(cityscapesDataSet(args.data_dir,
args.data_list,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
with torch.no_grad():
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, _, name = batch
#image = image.resize((1200,600),Image.BICUBIC)
output, _ = model(Variable(image).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (save_path, name))
output_col.save('%s/%s_color.png' %
(save_path, name.split('.')[0]))
print(save_path)
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
w, h = map(int, args.input_size.split(','))
input_size = (w, h)
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
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)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
device = torch.device("cuda" if not args.cpu else "cpu")
model = model.to(device)
model.eval()
testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, crop_size=input_size, mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
batch_size=1, shuffle=False, pin_memory=True)
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True)
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
image = image.to(device)
if args.model == 'DeeplabMulti':
output1, output2,_,_ = model(image)
output = interp(output2).cpu().data[0].numpy()
elif args.model == 'DeeplabVGG' or args.model == 'Oracle':
output = model(image)
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
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]))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'Deeplab':
model = Res_Deeplab(num_classes=args.num_classes)
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
### for running different versions of pytorch
model_dict = model.state_dict()
saved_state_dict = {k: v for k, v in saved_state_dict.items() if k in model_dict}
model_dict.update(saved_state_dict)
###
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, crop_size=(1024, 512), mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
batch_size=1, shuffle=False, pin_memory=True)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(1024, 2048), mode='bilinear', align_corners=True)
else:
interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
if args.model == 'Deeplab':
output = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
elif args.model == 'DeeplabVGG':
output = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
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]))
def eval(pth, cityscapes_eval_dir, i_iter):
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
if args.model == 'ResNet':
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(pth)
elif args.model == 'VGG':
model = DeeplabVGG(num_classes=args.num_classes)
saved_state_dict = torch.load(pth)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
cityscapesloader = data.DataLoader(cityscapesDataSet(
args.cityscapes_data_dir,
args.cityscapes_data_list,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
for index, batch in enumerate(cityscapesloader):
with torch.no_grad():
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
output = model(Variable(image).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (cityscapes_eval_dir, name))
output_col.save('%s/%s_color.png' %
(cityscapes_eval_dir, name.split('.')[0]))
if i_iter == 0:
break
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
for iter in range(ITER_START, ITER_END + 1, SPAN):
print('{0} /{1} processed'.format(iter, ITER_END))
if not os.path.exists(args.save.format(iter)):
os.makedirs(args.save.format(iter))
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from.format(iter))
for k, v in saved_state_dict.items():
if k.count('num_batches_tracked'):
del saved_state_dict[k]
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(cityscapesDataSet(args.data_dir,
args.data_list,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
for index, batch in enumerate(testloader):
if index % 100 == 0:
print '%d processd of %d' % (index, len(testloader))
image, _, name = batch
output1, output2 = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output2).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
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.format(iter), name.split('.')[0]))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
f = open(args.snapshot_dir + 'Evaluation.txt', 'w')
model = SEANet(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda()
testloader = data.DataLoader(cityscapesDataSet(args.data_dir,
args.data_list,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set='val'),
batch_size=1,
shuffle=False,
pin_memory=True)
input_size_target = (2048, 1024)
interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
test_mIoU(f, model, testloader, 0, input_size_target, print_per_batches=10)
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, _, name = batch
_, output = model(image.cuda())
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (args.snapshot_dir, name))
output_col.save('%s/%s_color.png' %
(args.snapshot_dir, name.split('.')[0]))
f.close()
def get_target_val_dataloader(args):
val_loader = DataLoader(cityscapesDataSet(args.root,
args.val.list,
crop_size=args.val.images_size,
mean=args.mean,
scale=False,
mirror=False,
set=args.val.set),
batch_size=args.batch_size,
shuffle=False,
pin_memory=True)
return val_loader
def show_val(seg_state_dict, save_dir, gpu_id):
"""Create the model and start the evaluation process."""
# args = get_arguments()
# save_dir=save_path.format(iter)
gpu0 = gpu_id
if not os.path.exists(save_dir):
os.makedirs(save_dir)
model = Res_Deeplab(num_classes=NUM_CLASSES)
model = nn.DataParallel(model, device_ids=device_ids)
model.load_state_dict(seg_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(cityscapesDataSet(DATA_DIRECTORY,
DATA_LIST_PATH,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=SET),
batch_size=1,
shuffle=False,
pin_memory=True)
# interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
# 作者建议：
interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
for index, batch in enumerate(testloader):
image, _, name = batch
with torch.no_grad():
output1, output2 = model(Variable(image).cuda(gpu0))
output = interp(output2).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
name = name[0].split('/')[-1]
output_col.save('%s/%s.png' % (save_dir, name.split('.')[0]))
print('colored pictures saving is done')
mIoUs = compute_mIoU(gtDir, save_dir, devkitDir)
return mIoUs
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
os.makedirs(args.save, exist_ok=True)
device = torch.device("cuda" if not args.cpu else "cpu")
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.to(device)
testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, crop_size=(1024, 512), mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
batch_size=1, shuffle=False, pin_memory=True)
interp = nn.Upsample(size=(1024, 2048), mode='bilinear', align_corners=True)
print('### STARTING EVALUATING ###')
print('total to process: %d' % len(testloader))
with torch.no_grad():
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processed' % index)
image, _, name, = batch
output1, output2 = model(image.to(device))
output = interp(output1 + output2).cpu().data[0].numpy()
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
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]))
print('### EVALUATING FINISHED ###')
def proceed_test(model, input_size, quick_test=1e10):
logger.info("proceed test on cityscapes val set...")
model.eval()
model.cuda()
testloader = data.DataLoader(cityscapesDataSet(crop_size=input_size,
mean=IMG_MEAN,
scale=False,
mirror=False,
set="val"),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(cityscape_image_size[1],
cityscape_image_size[0]),
mode='bilinear')
from tensorpack.utils.stats import MIoUStatistics
stat = MIoUStatistics(NUM_CLASSES)
for index, batch in tqdm(enumerate(testloader), desc="validation"):
if index > quick_test: break
image, label, _, name = batch
image, label = Variable(image, volatile=True), Variable(label)
output2 = model(image.cuda()) #(1,19,129,257)
output = interp(output2).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
stat.feed(output, label.data.cpu().numpy().squeeze())
miou16 = np.sum(stat.IoU) / 16
print("tensorpack class16 IoU with: {}".format(miou16))
model.train()
return miou16
def main():
"""Create the model and start the training."""
args = get_arguments()
if os.path.exists(args.snapshot_dir) == False:
os.mkdir(args.snapshot_dir)
f = open(args.snapshot_dir + 'GTA2Cityscapes_log.txt', 'w')
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)
# Create network
student_net = SEANet(num_classes=args.num_classes)
teacher_net = SEANet(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
new_params = student_net.state_dict().copy()
for i, j in zip(saved_state_dict, new_params):
if (i[0] != 'f') & (i[0] != 's') & (i[0] != 'u'):
new_params[j] = saved_state_dict[i]
student_net.load_state_dict(new_params)
teacher_net.load_state_dict(new_params)
for name, param in teacher_net.named_parameters():
param.requires_grad = False
teacher_net = teacher_net.cuda()
student_net = student_net.cuda()
src_loader = data.DataLoader(GTA5DataSet(args.data_dir_source,
args.data_list_source,
crop_size=input_size,
scale=False,
mirror=False,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
tgt_loader = data.DataLoader(cityscapesDataSet(args.data_dir_target,
args.data_list_target,
max_iters=24966,
crop_size=input_size,
scale=False,
mirror=False,
mean=IMG_MEAN,
set='val'),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
val_loader = data.DataLoader(cityscapesDataSet(args.data_dir_target,
args.data_list_target,
max_iters=None,
crop_size=input_size,
scale=False,
mirror=False,
mean=IMG_MEAN,
set='val'),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
num_batches = min(len(src_loader), len(tgt_loader))
optimizer = optim.Adam(student_net.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
optimizer.zero_grad()
student_params = list(student_net.parameters())
teacher_params = list(teacher_net.parameters())
teacher_optimizer = WeightEMA(
teacher_params,
student_params,
alpha=args.teacher_alpha,
)
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
n_class = args.num_classes
num_steps = args.num_epoch * num_batches
loss_hist = np.zeros((num_steps, 5))
index_i = -1
OA_hist = 0.2
aug_loss = torch.nn.MSELoss()
for epoch in range(args.num_epoch):
if epoch == 6:
return
for batch_index, (src_data,
tgt_data) in enumerate(zip(src_loader, tgt_loader)):
index_i += 1
tem_time = time.time()
student_net.train()
optimizer.zero_grad()
# train with source
images, src_label, _, im_name = src_data
images = images.cuda()
src_label = src_label.cuda()
_, src_output = student_net(images)
src_output = interp(src_output)
# Segmentation Loss
cls_loss_value = loss_calc(src_output, src_label)
_, predict_labels = torch.max(src_output, 1)
lbl_pred = predict_labels.detach().cpu().numpy()
lbl_true = src_label.detach().cpu().numpy()
metrics_batch = []
for lt, lp in zip(lbl_true, lbl_pred):
_, _, mean_iu, _ = label_accuracy_score(
lt, lp, n_class=args.num_classes)
metrics_batch.append(mean_iu)
miu = np.mean(metrics_batch, axis=0)
# train with target
images, label_target, _, im_name = tgt_data
images = images.cuda()
label_target = label_target.cuda()
tgt_t_input = images + torch.randn(
images.size()).cuda() * args.noise
tgt_s_input = images + torch.randn(
images.size()).cuda() * args.noise
_, tgt_s_output = student_net(tgt_s_input)
t_confidence, tgt_t_output = teacher_net(tgt_t_input)
t_confidence = t_confidence.squeeze()
# self-ensembling Loss
tgt_t_predicts = F.softmax(tgt_t_output,
dim=1).transpose(1, 2).transpose(2, 3)
tgt_s_predicts = F.softmax(tgt_s_output,
dim=1).transpose(1, 2).transpose(2, 3)
mask = t_confidence > args.attention_threshold
mask = mask.view(-1)
num_pixel = mask.shape[0]
mask_rate = torch.sum(mask).float() / num_pixel
tgt_s_predicts = tgt_s_predicts.contiguous().view(-1, n_class)
tgt_s_predicts = tgt_s_predicts[mask]
tgt_t_predicts = tgt_t_predicts.contiguous().view(-1, n_class)
tgt_t_predicts = tgt_t_predicts[mask]
aug_loss_value = aug_loss(tgt_s_predicts, tgt_t_predicts)
aug_loss_value = args.st_weight * aug_loss_value
# TOTAL LOSS
if mask_rate == 0.:
aug_loss_value = torch.tensor(0.).cuda()
total_loss = cls_loss_value + aug_loss_value
total_loss.backward()
loss_hist[index_i, 0] = total_loss.item()
loss_hist[index_i, 1] = cls_loss_value.item()
loss_hist[index_i, 2] = aug_loss_value.item()
loss_hist[index_i, 3] = miu
optimizer.step()
teacher_optimizer.step()
batch_time = time.time() - tem_time
if (batch_index + 1) % 10 == 0:
print(
'epoch %d/%d: %d/%d time: %.2f miu = %.1f cls_loss = %.3f st_loss = %.3f \n'
% (epoch + 1, args.num_epoch, batch_index + 1, num_batches,
batch_time,
np.mean(loss_hist[index_i - 9:index_i + 1, 3]) * 100,
np.mean(loss_hist[index_i - 9:index_i + 1, 1]),
np.mean(loss_hist[index_i - 9:index_i + 1, 2])))
f.write(
'epoch %d/%d: %d/%d time: %.2f miu = %.1f cls_loss = %.3f st_loss = %.3f \n'
% (epoch + 1, args.num_epoch, batch_index + 1, num_batches,
batch_time,
np.mean(loss_hist[index_i - 9:index_i + 1, 3]) * 100,
np.mean(loss_hist[index_i - 9:index_i + 1, 1]),
np.mean(loss_hist[index_i - 9:index_i + 1, 2])))
f.flush()
if (batch_index + 1) % 500 == 0:
OA_new = test_mIoU(f,
teacher_net,
val_loader,
epoch + 1,
input_size_target,
print_per_batches=10)
# Saving the models
if OA_new > OA_hist:
f.write('Save Model\n')
print('Save Model')
model_name = 'GTA2Cityscapes_epoch' + repr(
epoch + 1) + 'batch' + repr(batch_index +
1) + 'tgt_miu_' + repr(
int(OA_new *
1000)) + '.pth'
torch.save(teacher_net.state_dict(),
os.path.join(args.snapshot_dir, model_name))
OA_hist = OA_new
f.close()
torch.save(teacher_net.state_dict(),
os.path.join(args.snapshot_dir, 'GTA_TeacherNet.pth'))
np.savez(args.snapshot_dir + 'GTA_loss.npz', loss_hist=loss_hist)
def main():
"""Create the model and start the training."""
w, h = map(int, args.input_size.split(','))
args.input_size = (w, h)
w, h = map(int, args.crop_size.split(','))
args.crop_size = (h, w)
w, h = map(int, args.input_size_target.split(','))
args.input_size_target = (w, h)
cudnn.enabled = True
cudnn.benchmark = True
str_ids = args.gpu_ids.split(',')
gpu_ids = []
for str_id in str_ids:
gid = int(str_id)
if gid >=0:
gpu_ids.append(gid)
num_gpu = len(gpu_ids)
args.multi_gpu = False
if num_gpu>1:
args.multi_gpu = True
Trainer = AD_Trainer(args)
Trainer.G = torch.nn.DataParallel( Trainer.G, gpu_ids)
Trainer.D1 = torch.nn.DataParallel( Trainer.D1, gpu_ids)
Trainer.D2 = torch.nn.DataParallel( Trainer.D2, gpu_ids)
else:
Trainer = AD_Trainer(args)
print(Trainer)
trainloader = data.DataLoader(
cityscapes_pseudo_DataSet(args.data_dir, args.data_list, max_iters=args.num_steps * args.iter_size * args.batch_size,
resize_size=args.input_size,
crop_size=args.crop_size,
scale=True, mirror=True, mean=IMG_MEAN,
set='train', autoaug = args.autoaug, synthia=True),
batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True, drop_last=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,
resize_size=args.input_size_target,
crop_size=args.crop_size,
scale=False, mirror=args.random_mirror, mean=IMG_MEAN,
set=args.set, autoaug = args.autoaug_target),
batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers,
pin_memory=True, drop_last=True)
targetloader_iter = enumerate(targetloader)
# set up tensor board
if args.tensorboard:
args.log_dir += '/'+ os.path.basename(args.snapshot_dir)
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
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
adjust_learning_rate(Trainer.gen_opt , i_iter, args)
#adjust_learning_rate_D(Trainer.dis1_opt, i_iter, args)
#adjust_learning_rate_D(Trainer.dis2_opt, i_iter, args)
for sub_i in range(args.iter_size):
# train G
# train with source
_, batch = trainloader_iter.__next__()
_, batch_t = targetloader_iter.__next__()
images, labels, _, _ = batch
images = images.cuda()
labels = labels.long().cuda()
images_t, labels_t, _, _ = batch_t
images_t = images_t.cuda()
labels_t = labels_t.long().cuda()
with Timer("Elapsed time in update: %f"):
loss_seg1, loss_seg2, loss_adv_target1, loss_adv_target2, loss_me, loss_kl, pred1, pred2, pred_target1, pred_target2, val_loss = Trainer.gen_update(images, images_t, labels, labels_t, i_iter)
loss_seg_value1 += loss_seg1.item() / args.iter_size
loss_seg_value2 += loss_seg2.item() / args.iter_size
loss_adv_target_value1 += loss_adv_target1 / args.iter_size
loss_adv_target_value2 += loss_adv_target2 / args.iter_size
loss_me_value = loss_me
if args.lambda_adv_target1 > 0 and args.lambda_adv_target2 > 0:
loss_D1, loss_D2 = Trainer.dis_update(pred1, pred2, pred_target1, pred_target2)
loss_D_value1 += loss_D1.item()
loss_D_value2 += loss_D2.item()
else:
loss_D_value1 = 0
loss_D_value2 = 0
del pred1, pred2, pred_target1, pred_target2
if args.tensorboard:
scalar_info = {
'loss_seg1': loss_seg_value1,
'loss_seg2': loss_seg_value2,
'loss_adv_target1': loss_adv_target_value1,
'loss_adv_target2': loss_adv_target_value2,
'loss_me_target': loss_me_value,
'loss_kl_target': loss_kl,
'loss_D1': loss_D_value1,
'loss_D2': loss_D_value2,
'val_loss': val_loss,
}
if i_iter % 100 == 0:
for key, val in scalar_info.items():
writer.add_scalar(key, val, i_iter)
print('exp = {}'.format(args.snapshot_dir))
print(
'\033[1m iter = %8d/%8d \033[0m loss_seg1 = %.3f loss_seg2 = %.3f loss_me = %.3f loss_kl = %.3f loss_adv1 = %.3f, loss_adv2 = %.3f loss_D1 = %.3f loss_D2 = %.3f, val_loss=%.3f'%(i_iter, args.num_steps, loss_seg_value1, loss_seg_value2, loss_me_value, loss_kl, loss_adv_target_value1, loss_adv_target_value2, loss_D_value1, loss_D_value2, val_loss))
# clear loss
del loss_seg1, loss_seg2, loss_adv_target1, loss_adv_target2, loss_me, loss_kl, val_loss
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(Trainer.G.state_dict(), osp.join(args.snapshot_dir, 'GTA5_' + str(args.num_steps_stop) + '.pth'))
torch.save(Trainer.D1.state_dict(), osp.join(args.snapshot_dir, 'GTA5_' + str(args.num_steps_stop) + '_D1.pth'))
torch.save(Trainer.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(Trainer.G.state_dict(), osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth'))
torch.save(Trainer.D1.state_dict(), osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D1.pth'))
torch.save(Trainer.D2.state_dict(), osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D2.pth'))
if args.tensorboard:
writer.close()
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
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)
if CONTINUE_FLAG==1:
model.load_state_dict(saved_state_dict)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# init D
model_D_a = FCDiscriminator(num_classes=256) # need to check
model_D1 = FCDiscriminator(num_classes=args.num_classes)
model_D2 = FCDiscriminator(num_classes=args.num_classes)
if CONTINUE_FLAG==1:
d1_saved_state_dict = torch.load(D1_RESTORE_FROM)
d2_saved_state_dict = torch.load(D2_RESTORE_FROM)
model_D1.load_state_dict(d1_saved_state_dict)
model_D2.load_state_dict(d2_saved_state_dict)
model_D_a.train()
model_D_a.cuda(args.gpu)
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=False, 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_D_a = optim.Adam(model_D_a.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
optimizer_D_a.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
mIoUs = []
for i_iter in range(continue_start_iter+1,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_D_a.zero_grad()
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate_D(optimizer_D_a, 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_D_a.parameters():
param.requires_grad = False
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)
pred_a, pred1, pred2 = model(images)
pred1=nn.functional.interpolate(pred1,size=(input_size[1], input_size[0]), mode='bilinear',align_corners=True)
pred2 = nn.functional.interpolate(pred2, size=(input_size[1], input_size[0]), mode='bilinear',
align_corners=True)
# 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)
lambda_wtight = (80000 - i_iter) / 80000
if lambda_wtight > 0:
pred_target_a, _, _ = model(images)
D_out_a = model_D_a(pred_target_a)
loss_adv_target_a = bce_loss(D_out_a,
Variable(torch.FloatTensor(D_out_a.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_adv_target_a=LAMBDA_ADV_TARGET_A *loss_adv_target_a
loss_adv_target_a = loss_adv_target_a / args.iter_size
loss_adv_target_a.backward()
_, pred_target1, pred_target2 = model(images)
pred_target1 = nn.functional.interpolate(pred_target1,size=(input_size_target[1], input_size_target[0]), mode='bilinear',align_corners=True)
pred_target2 = nn.functional.interpolate(pred_target2, size=(input_size_target[1], input_size_target[0]),
mode='bilinear', align_corners=True)
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_D_a.parameters():
param.requires_grad = True
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
# train with source
if lambda_wtight > 0:
pred_a=pred_a.detach()
D_out_a = model_D_a(pred_a)
loss_D_a = bce_loss(D_out_a,
Variable(torch.FloatTensor(D_out_a.data.size()).fill_(source_label)).cuda(args.gpu))
loss_D_a = loss_D_a / args.iter_size / 2
loss_D_a.backward()
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_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
if lambda_wtight > 0:
pred_target_a=pred_target_a.detach()
D_out_a = model_D_a(pred_target_a)
loss_D_a = bce_loss(D_out_a,
Variable(torch.FloatTensor(D_out_a.data.size()).fill_(target_label)).cuda(args.gpu))
loss_D_a = loss_D_a / args.iter_size / 2
loss_D_a.backward()
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))
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_D_a.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'))
show_val(model.state_dict(), LAMBDA_ADV_TARGET_A ,i_iter)
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'))
mIoU=show_val(model.state_dict(), LAMBDA_ADV_TARGET_A ,i_iter)
mIoUs.append(str(round(np.nanmean(mIoU) * 100, 2)))
for miou in mIoUs:
print(miou)
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
if not os.path.exists(args.save):
os.makedirs(args.save)
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.eval()
model.cuda(gpu0)
if os.path.isfile(citys_feat_distr_path) == False:
testloader = data.DataLoader(cityscapesDataSet(args.data_dir,
args.data_list,
crop_size=(1024, 512),
mean=CITY_IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=1,
shuffle=False,
pin_memory=True)
# interp = nn.Upsample(size=(1024, 2048), mode='bilinear', align_corners=True)
interp_down = nn.Upsample(size=(16, 32),
mode='bilinear',
align_corners=True)
citys_feat_distrs = []
citys_img_paths = []
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd of %d' % (index, len(testloader)))
image, _, name = batch
output1, output2 = model(Variable(image, volatile=True).cuda(gpu0))
output = interp_down(output2).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = output[np.newaxis, :] # add a dim
citys_feat_distrs.extend(output)
citys_img_paths.extend(name)
#name: 'frankfurt/frankfurt_000001_007973_leftImg8bit.png'
# name = name[0].split('/')[-1]
citys_feat_distrs_np = np.array(citys_feat_distrs)
citys_img_paths_np = np.array(citys_img_paths)
np.save(citys_feat_distr_path, citys_feat_distrs_np)
np.save(citys_imgpaths_path, citys_img_paths_np)
else:
citys_feat_distrs_np = np.load(citys_feat_distr_path)
citys_img_paths_np = np.load(citys_imgpaths_path)
if os.path.isfile(gta_feat_distr_path) == False:
gtaloader = data.DataLoader(GTA5DataSet(GTA_DATA_DIRECTORY,
GTA_DATA_LIST_PATH,
crop_size=(1024, 512),
mean=GTA_IMG_MEAN,
scale=False,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
interp_down = nn.Upsample(size=(16, 32),
mode='bilinear',
align_corners=True)
gta_feat_distrs = []
gta_img_paths = []
for index, batch in enumerate(gtaloader):
if index % 100 == 0:
print('%d processd of %d' % (index, len(gtaloader)))
image, _, _, name = batch
output1, output2 = model(Variable(image, volatile=True).cuda(gpu0))
output = interp_down(output2).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = output[np.newaxis, :] # add a dim
gta_feat_distrs.extend(output)
gta_img_paths.extend(name)
gta_feat_distrs_np = np.array(gta_feat_distrs)
gta_img_paths_np = np.array(gta_img_paths)
np.save(gta_feat_distr_path, gta_feat_distrs_np)
np.save(gta_imgpaths_path, gta_img_paths_np)
else:
gta_feat_distrs_np = np.load(gta_feat_distr_path)
gta_img_paths_np = np.load(gta_imgpaths_path)
if os.path.isfile(closest_imgs_path) == False:
temp_feat = citys_feat_distrs_np[0, :]
# [m,n,c]=temp_feat.shape
pixel_amount = temp_feat.size
closest_imgs_locs = []
for i in range(citys_img_paths_np.shape[0]):
cur_citys_feat = citys_feat_distrs_np[i, :]
distances = []
if i % 10 == 0:
print(i)
for j in range(gta_img_paths_np.shape[0]):
cur_gta_feat = gta_feat_distrs_np[j, :]
dist_abs = abs(cur_citys_feat - cur_gta_feat)
# e_dist = np.sqrt(np.square(dist_abs).sum(axis=1))
dist_mean = np.sum(dist_abs) / pixel_amount
distances.append(dist_mean)
min_loc = np.argsort(distances)
# need to check overlap
top_ord = 3
closest_imgs_loc = min_loc[:top_ord]
intersect_imgs = np.intersect1d(closest_imgs_loc,
closest_imgs_locs)
while intersect_imgs.size:
inters_num = len(intersect_imgs)
closest_imgs_loc_confirm = np.setdiff1d(
closest_imgs_loc, intersect_imgs) # find the difference
closest_imgs_loc_candi = min_loc[top_ord:top_ord + inters_num]
top_ord = top_ord + inters_num
closest_imgs_loc_confirm = np.concatenate(
[closest_imgs_loc_confirm, closest_imgs_loc_candi])
closest_imgs_loc = closest_imgs_loc_confirm
intersect_imgs = np.intersect1d(closest_imgs_loc,
closest_imgs_locs)
closest_imgs_locs.extend(closest_imgs_loc)
np.save(closest_imgs_path, closest_imgs_locs)
else:
closest_imgs_locs = np.load(closest_imgs_path)
closest_imgs_locs_uni = np.unique(closest_imgs_locs)
zq = 1
# get file_names
with open(src_train_imgs_txt, 'w') as f_train:
for img_num in closest_imgs_locs_uni:
line = gta_img_paths_np[img_num] + '\n'
f_train.write(line)
def main():
"""Create the model and start the training."""
device = torch.device("cuda" if not args.cpu else "cpu")
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)
cudnn.enabled = True
# Create network
if args.model == 'ResNet':
model = DeeplabMulti(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 == 'VGG':
model = DeeplabVGG(num_classes=args.num_classes,
vgg16_caffe_path='./model/vgg16_init.pth',
pretrained=True)
model.train()
model.to(device)
cudnn.benchmark = True
# init D
if args.model == 'ResNet':
model_D = FCDiscriminator(num_classes=2048).to(device)
if args.model == 'VGG':
model_D = FCDiscriminator(num_classes=1024).to(device)
model_D.train()
model_D.to(device)
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_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
seg_loss = torch.nn.CrossEntropyLoss(ignore_index=255)
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
# set up tensor board
if args.tensorboard:
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
for i_iter in range(args.num_steps):
loss_seg = 0
loss_adv_target_value = 0
loss_D_value = 0
loss_cla_value = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
# train G
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# train with source
_, batch = trainloader_iter.__next__()
images, labels, _, _ = batch
images = images.to(device)
labels = labels.long().to(device)
feature, prediction = model(images)
prediction = interp(prediction)
loss = seg_loss(prediction, labels)
loss.backward()
loss_seg = loss.item()
# train with target
_, batch = targetloader_iter.__next__()
images, _, _ = batch
images = images.to(device)
feature_target, _ = model(images)
_, D_out = model_D(feature_target)
loss_adv_target = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(source_label).to(device))
#print(args.lambda_adv_target)
loss = args.lambda_adv_target * loss_adv_target
loss.backward()
loss_adv_target_value = loss_adv_target.item()
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with source
feature = feature.detach()
cla, D_out = model_D(feature)
cla = interp(cla)
loss_cla = seg_loss(cla, labels)
loss_D = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(source_label).to(device))
loss_D = loss_D / 2
#print(args.lambda_s)
loss_Disc = args.lambda_s * loss_cla + loss_D
loss_Disc.backward()
loss_cla_value = loss_cla.item()
loss_D_value = loss_D.item()
# train with target
feature_target = feature_target.detach()
_, D_out = model_D(feature_target)
loss_D = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(target_label).to(device))
loss_D = loss_D / 2
loss_D.backward()
loss_D_value += loss_D.item()
optimizer.step()
optimizer_D.step()
if args.tensorboard:
scalar_info = {
'loss_seg': loss_seg,
'loss_cla': loss_cla_value,
'loss_adv_target': loss_adv_target_value,
'loss_D': loss_D_value,
}
if i_iter % 10 == 0:
for key, val in scalar_info.items():
writer.add_scalar(key, val, i_iter)
#print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f} loss_adv = {3:.3f} loss_D = {4:.3f} loss_cla = {5:.3f}'
.format(i_iter, args.num_steps, loss_seg, loss_adv_target_value,
loss_D_value, loss_cla_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_D.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '_D.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_D.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D.pth'))
if args.tensorboard:
writer.close()
def main():
"""Create the model and start the training."""
device = torch.device("cuda" if not args.cpu else "cpu")
cudnn.benchmark = True
cudnn.enabled = True
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)
Iter = 0
bestIoU = 0
# Create network
# init G
if args.model == 'DeepLab':
model = DeeplabMultiFeature(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)
if args.continue_train:
if list(saved_state_dict.keys())[0].split('.')[0] == 'module':
for key in saved_state_dict.keys():
saved_state_dict['.'.join(
key.split('.')[1:])] = saved_state_dict.pop(key)
model.load_state_dict(saved_state_dict)
else:
new_params = model.state_dict().copy()
for i in saved_state_dict:
i_parts = i.split('.')
if not args.num_classes == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
model.load_state_dict(new_params)
# init D
model_D = FCDiscriminator(num_classes=args.num_classes).to(device)
if args.continue_train:
model_weights_path = args.restore_from
temp = model_weights_path.split('.')
temp[-2] = temp[-2] + '_D'
model_D_weights_path = '.'.join(temp)
model_D.load_state_dict(torch.load(model_D_weights_path))
temp = model_weights_path.split('.')
temp = temp[-2][-9:]
Iter = int(temp.split('_')[1]) + 1
model.train()
model.to(device)
model_D.train()
model_D.to(device)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
# init data loader
if args.data_dir.split('/')[-1] == 'gta5_deeplab':
trainset = 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)
elif args.data_dir.split('/')[-1] == 'syn_deeplab':
trainset = 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)
trainloader = data.DataLoader(trainset,
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)
# init optimizer
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
model, optimizer = amp.initialize(model,
optimizer,
opt_level="O2",
keep_batchnorm_fp32=True,
loss_scale="dynamic")
model_D, optimizer_D = amp.initialize(model_D,
optimizer_D,
opt_level="O2",
keep_batchnorm_fp32=True,
loss_scale="dynamic")
# init loss
bce_loss = torch.nn.BCEWithLogitsLoss()
seg_loss = torch.nn.CrossEntropyLoss(ignore_index=255)
L1_loss = torch.nn.L1Loss(reduction='none')
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)
test_interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
# labels for adversarial training
source_label = 0
target_label = 1
# init prototype
num_prototype = args.num_prototype
num_ins = args.num_prototype * 10
src_cls_features = torch.zeros([len(BG_LABEL), num_prototype, 2048],
dtype=torch.float32).to(device)
src_cls_ptr = np.zeros(len(BG_LABEL), dtype=np.uint64)
src_ins_features = torch.zeros([len(FG_LABEL), num_ins, 2048],
dtype=torch.float32).to(device)
src_ins_ptr = np.zeros(len(FG_LABEL), dtype=np.uint64)
# set up tensor board
if args.tensorboard:
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
# start training
for i_iter in range(Iter, args.num_steps):
loss_seg_value = 0
loss_adv_target_value = 0
loss_D_value = 0
loss_cls_value = 0
loss_ins_value = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# train with source
_, batch = trainloader_iter.__next__()
images, labels, _, _ = batch
images = images.to(device)
labels = labels.long().to(device)
src_feature, pred = model(images)
pred_softmax = F.softmax(pred, dim=1)
pred_idx = torch.argmax(pred_softmax, dim=1)
right_label = F.interpolate(labels.unsqueeze(0).float(),
(pred_idx.size(1), pred_idx.size(2)),
mode='nearest').squeeze(0).long()
right_label[right_label != pred_idx] = 255
for ii in range(len(BG_LABEL)):
cls_idx = BG_LABEL[ii]
mask = right_label == cls_idx
if torch.sum(mask) == 0:
continue
feature = global_avg_pool(src_feature, mask.float())
if cls_idx != torch.argmax(
torch.squeeze(model.layer6(
feature.half()).float())).item():
continue
src_cls_features[ii,
int(src_cls_ptr[ii] %
num_prototype), :] = torch.squeeze(
feature).clone().detach()
src_cls_ptr[ii] += 1
seg_ins = seg_label(right_label.squeeze())
for ii in range(len(FG_LABEL)):
cls_idx = FG_LABEL[ii]
segmask, pixelnum = seg_ins[ii]
if len(pixelnum) == 0:
continue
sortmax = np.argsort(pixelnum)[::-1]
for i in range(min(10, len(sortmax))):
mask = segmask == (sortmax[i] + 1)
feature = global_avg_pool(src_feature, mask.float())
if cls_idx != torch.argmax(
torch.squeeze(
model.layer6(feature.half()).float())).item():
continue
src_ins_features[ii, int(src_ins_ptr[ii] %
num_ins), :] = torch.squeeze(
feature).clone().detach()
src_ins_ptr[ii] += 1
pred = interp(pred)
loss_seg = seg_loss(pred, labels)
loss = loss_seg
# proper normalization
loss = loss / args.iter_size
amp_backward(loss, optimizer)
loss_seg_value += loss_seg.item() / args.iter_size
# train with target
_, batch = targetloader_iter.__next__()
images, _, _ = batch
images = images.to(device)
trg_feature, pred_target = model(images)
pred_target_softmax = F.softmax(pred_target, dim=1)
pred_target_idx = torch.argmax(pred_target_softmax, dim=1)
loss_cls = torch.zeros(1).to(device)
loss_ins = torch.zeros(1).to(device)
if i_iter > 0:
for ii in range(len(BG_LABEL)):
cls_idx = BG_LABEL[ii]
if src_cls_ptr[ii] / num_prototype <= 1:
continue
mask = pred_target_idx == cls_idx
feature = global_avg_pool(trg_feature, mask.float())
if cls_idx != torch.argmax(
torch.squeeze(
model.layer6(feature.half()).float())).item():
continue
ext_feature = feature.squeeze().expand(num_prototype, 2048)
loss_cls += torch.min(
torch.sum(L1_loss(ext_feature,
src_cls_features[ii, :, :]),
dim=1) / 2048.)
seg_ins = seg_label(pred_target_idx.squeeze())
for ii in range(len(FG_LABEL)):
cls_idx = FG_LABEL[ii]
if src_ins_ptr[ii] / num_ins <= 1:
continue
segmask, pixelnum = seg_ins[ii]
if len(pixelnum) == 0:
continue
sortmax = np.argsort(pixelnum)[::-1]
for i in range(min(10, len(sortmax))):
mask = segmask == (sortmax[i] + 1)
feature = global_avg_pool(trg_feature, mask.float())
feature = feature.squeeze().expand(num_ins, 2048)
loss_ins += torch.min(
torch.sum(L1_loss(feature,
src_ins_features[ii, :, :]),
dim=1) / 2048.) / min(10, len(sortmax))
pred_target = interp_target(pred_target)
D_out = model_D(F.softmax(pred_target, dim=1))
loss_adv_target = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(source_label).to(device))
loss = args.lambda_adv_target * loss_adv_target + args.lambda_adv_cls * loss_cls + args.lambda_adv_ins * loss_ins
loss = loss / args.iter_size
amp_backward(loss, optimizer)
loss_adv_target_value += loss_adv_target.item() / args.iter_size
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with source
pred = pred.detach()
D_out = model_D(F.softmax(pred, dim=1))
loss_D = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(source_label).to(device))
loss_D = loss_D / args.iter_size / 2
amp_backward(loss_D, optimizer_D)
loss_D_value += loss_D.item()
# train with target
pred_target = pred_target.detach()
D_out = model_D(F.softmax(pred_target, dim=1))
loss_D = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(target_label).to(device))
loss_D = loss_D / args.iter_size / 2
amp_backward(loss_D, optimizer_D)
loss_D_value += loss_D.item()
optimizer.step()
optimizer_D.step()
if args.tensorboard:
scalar_info = {
'loss_seg': loss_seg_value,
'loss_adv_target': loss_adv_target_value,
'loss_D': loss_D_value,
}
if i_iter % 10 == 0:
for key, val in scalar_info.items():
writer.add_scalar(key, val, i_iter)
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv = {3:.3f} loss_D = {4:.3f} loss_cls = {5:.3f} loss_ins = {6:.3f}'
.format(i_iter, args.num_steps, loss_seg_value,
loss_adv_target_value, loss_D_value, loss_cls.item(),
loss_ins.item()))
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_D.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '_D.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
if not os.path.exists(args.save):
os.makedirs(args.save)
testloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target_test,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set='val'),
batch_size=1,
shuffle=False,
pin_memory=True)
model.eval()
for index, batch in enumerate(testloader):
image, _, name = batch
with torch.no_grad():
output1, output2 = model(Variable(image).to(device))
output = test_interp(output2).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (args.save, name))
mIoUs = compute_mIoU(osp.join(args.data_dir_target, 'gtFine/val'),
args.save, 'dataset/cityscapes_list')
mIoU = round(np.nanmean(mIoUs) * 100, 2)
if mIoU > bestIoU:
bestIoU = mIoU
torch.save(model.state_dict(),
osp.join(args.snapshot_dir, 'BestGTA5.pth'))
torch.save(model_D.state_dict(),
osp.join(args.snapshot_dir, 'BestGTA5_D.pth'))
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D.pth'))
model.train()
if args.tensorboard:
writer.close()
def main():
"""Create the model and start the training."""
device = torch.device("cuda" if not args.cpu else "cpu")
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)
cudnn.enabled = True
# Create network
if args.model == 'DeepLab':
model = DeeplabMulti(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.to(device)
cudnn.benchmark = True
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes).to(device)
model_D2 = FCDiscriminator(num_classes=args.num_classes).to(device)
model_D1.train()
model_D1.to(device)
model_D2.train()
model_D2.to(device)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(cityscapestrainDataSet(
args.data_dir,
args.data_list,
args.data_label_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,
set=args.set),
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()
seg_loss = torch.nn.CrossEntropyLoss(ignore_index=255)
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
# set up tensor board
if args.tensorboard:
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
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 = images.to(device)
labels = labels.long().to(device)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
loss_seg1 = seg_loss(pred1, labels)
loss_seg2 = seg_loss(pred2, labels)
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 = targetloader_iter.__next__()
images, _, _ = batch
images = images.to(device)
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,
torch.FloatTensor(
D_out1.data.size()).fill_(source_label).to(device))
loss_adv_target2 = bce_loss(
D_out2,
torch.FloatTensor(
D_out2.data.size()).fill_(source_label).to(device))
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.item() / args.iter_size
loss_adv_target_value2 += loss_adv_target2.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,
torch.FloatTensor(
D_out1.data.size()).fill_(source_label).to(device))
loss_D2 = bce_loss(
D_out2,
torch.FloatTensor(
D_out2.data.size()).fill_(source_label).to(device))
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.item()
loss_D_value2 += loss_D2.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,
torch.FloatTensor(
D_out1.data.size()).fill_(target_label).to(device))
loss_D2 = bce_loss(
D_out2,
torch.FloatTensor(
D_out2.data.size()).fill_(target_label).to(device))
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.item()
loss_D_value2 += loss_D2.item()
optimizer.step()
optimizer_D1.step()
optimizer_D2.step()
if args.tensorboard:
scalar_info = {
'loss_seg1': loss_seg_value1,
'loss_seg2': loss_seg_value2,
'loss_adv_target1': loss_adv_target_value1,
'loss_adv_target2': loss_adv_target_value2,
'loss_D1': loss_D_value1,
'loss_D2': loss_D_value2,
}
if i_iter % 10 == 0:
for key, val in scalar_info.items():
writer.add_scalar(key, val, i_iter)
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_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'))
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'))
if args.tensorboard:
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)
cudnn.enabled = True
gpu = args.gpu
# Create network
if args.model == 'DeepLab':
model = DeeplabMulti(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
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)
# Implemented by Bongjoon Hyun
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)
#
# Implemented by Bongjoon Hyun
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.gan == 'Vanilla':
bce_loss = torch.nn.BCEWithLogitsLoss()
elif args.gan == 'LS':
# Implemented by Bongjoon Hyun
bce_loss = torch.nn.MSELoss()
#
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
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):
# Implemented by Bongjoon Hyun
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) / 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
_, 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)
D1_out = model_D1(F.softmax(pred_target1))
D2_out = model_D2(F.softmax(pred_target2))
labels_source1 = Variable(
torch.FloatTensor(
D1_out.data.size()).fill_(source_label)).cuda(args.gpu)
labels_source2 = Variable(
torch.FloatTensor(
D2_out.data.size()).fill_(source_label)).cuda(args.gpu)
loss_adv_target1 = bce_loss(D1_out, labels_source1)
loss_adv_target2 = bce_loss(D2_out, labels_source2)
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
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
pred1 = pred1.detach()
pred2 = pred2.detach()
D1_out = model_D1(F.softmax(pred1))
D2_out = model_D2(F.softmax(pred2))
labels_source1 = Variable(
torch.FloatTensor(
D1_out.data.size()).fill_(source_label)).cuda(args.gpu)
labels_source2 = Variable(
torch.FloatTensor(
D2_out.data.size()).fill_(source_label)).cuda(args.gpu)
loss_D1 = bce_loss(D1_out, labels_source1) / args.iter_size / 2
loss_D2 = bce_loss(D2_out, labels_source2) / 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()
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
D1_out = model_D1(F.softmax(pred_target1))
D2_out = model_D2(F.softmax(pred_target2))
labels_target1 = Variable(
torch.FloatTensor(
D1_out.data.size()).fill_(target_label)).cuda(args.gpu)
labels_target2 = Variable(
torch.FloatTensor(
D2_out.data.size()).fill_(target_label)).cuda(args.gpu)
loss_D1 = bce_loss(D1_out, labels_target1) / args.iter_size / 2
loss_D2 = bce_loss(D2_out, labels_target2) / 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_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'))
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.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."""
gpu_id_2 = 3
gpu_id_1 = 2
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)
cudnn.enabled = True
# gpu = args.gpu
# Create network
if args.model == 'DeepLab':
model = DeeplabMulti(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
print("from url")
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
print("from restore")
saved_state_dict = torch.load(args.restore_from)
saved_state_dict = torch.load(
'snapshots/GTA2Cityscapes_multi_54/GTA5_10000.pth')
model.load_state_dict(saved_state_dict)
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(gpu_id_2)
cudnn.benchmark = True
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes)
#model_D2 = model_D1
model_D2 = FCDiscriminator(num_classes=args.num_classes)
model_D1.train()
model_D1.cuda(gpu_id_1)
model_D2.train()
model_D2.cuda(gpu_id_1)
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)
_, batch_last = trainloader_iter.next()
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)
# print(args.num_steps * args.iter_size * args.batch_size, trainloader.__len__())
targetloader_iter = enumerate(targetloader)
_, batch_last_target = targetloader_iter.next()
# for i in range(200):
# _, batch = targetloader_iter.__next__()
# exit()
# 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.MSELoss()
def upsample_(input_):
return nn.functional.interpolate(input_,
size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=False)
def upsample_target(input_):
return nn.functional.interpolate(input_,
size=(input_size_target[1],
input_size_target[0]),
mode='bilinear',
align_corners=False)
interp = upsample_
interp_target = upsample_target
# labels for adversarial training
source_label = 1
target_label = -1
mix_label = 0
for i_iter in range(10000, 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
def result_model(batch, interp_):
images, labels, _, name = batch
images = Variable(images).cuda(gpu_id_2)
labels = Variable(labels.long()).cuda(gpu_id_1)
pred1, pred2 = model(images)
pred1 = interp_(pred1)
pred2 = interp_(pred2)
pred1_ = pred1.cuda(gpu_id_1)
pred2_ = pred2.cuda(gpu_id_1)
return pred1_, pred2_, labels
# train with source
# _, batch = trainloader_iter.next()
_, batch = trainloader_iter.next()
_, batch_target = targetloader_iter.next()
pred1, pred2, labels = result_model(batch, interp)
loss_seg1 = loss_calc(pred1, labels, gpu_id_1)
loss_seg2 = loss_calc(pred2, labels, gpu_id_1)
loss = loss_seg2 + args.lambda_seg * loss_seg1
loss = loss / args.iter_size
loss.backward()
loss_seg_1 = loss_seg1.data.cpu().numpy() / args.iter_size
loss_seg_2 = loss_seg2.data.cpu().numpy() / args.iter_size
# print(loss_seg_1, loss_seg_2)
pred1, pred2, labels = result_model(batch_target, interp_target)
loss_seg1 = loss_calc(pred1, labels, gpu_id_1)
loss_seg2 = loss_calc(pred2, labels, gpu_id_1)
loss = loss_seg2 + args.lambda_seg * loss_seg1
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
# output = pred2.cpu().data[0].numpy()
# real_lab = labels.cpu().data[0].numpy()
# output = output.transpose(1,2,0)
# print(real_lab.shape, output.shape)
# output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
# output_col = colorize_mask(output)
# real_lab_col = colorize_mask(real_lab)
# output = Image.fromarray(output)
# # name[0].split('/')[-1]
# # print('result/train_seg_result/' + name[0][len(name[0])-23:len(name[0])-4] + '_color.png')
# output_col.save('result/train_seg_result/' + name[0].split('/')[-1] + '_color.png')
# real_lab_col.save('result/train_seg_result/' + name[0].split('/')[-1] + '_real.png')
# print(loss_seg_value1, loss_seg_value2)
# if i_iter == 100:
# exit()
# else:
# break
# train with target
#_, batch = targetloader_iter.next()
# images, _, _ = target_batch
# images_target = Variable(images_target).cuda(gpu_id_2)
pred1_last_target, pred2_last_target, labels_last_target = result_model(
batch_last_target, interp_target)
pred1_target, pred2_target, labels_target = result_model(
batch_target, interp_target)
pred1_target_D = F.softmax((pred1_target), dim=1)
pred2_target_D = F.softmax((pred2_target), dim=1)
pred1_last_target_D = F.softmax((pred1_last_target), dim=1)
pred2_last_target_D = F.softmax((pred2_last_target), dim=1)
fake1_D = torch.cat((pred1_target_D, pred1_last_target_D), dim=1)
fake2_D = torch.cat((pred2_target_D, pred2_last_target_D), dim=1)
D_out_fake_1 = model_D1(fake1_D)
D_out_fake_2 = model_D2(fake1_D)
loss_adv_fake1 = bce_loss(
D_out_fake_1,
Variable(
torch.FloatTensor(D_out_fake_1.data.size()).fill_(
source_label)).cuda(gpu_id_1))
loss_adv_fake2 = bce_loss(
D_out_fake_2,
Variable(
torch.FloatTensor(D_out_fake_2.data.size()).fill_(
source_label)).cuda(gpu_id_1))
loss_adv_target1 = loss_adv_fake1
loss_adv_target2 = loss_adv_fake2
loss = args.lambda_adv_target1 * loss_adv_target1.cuda(
gpu_id_1) + args.lambda_adv_target2 * loss_adv_target2.cuda(
gpu_id_1)
loss = loss / args.iter_size
loss.backward()
pred1, pred2, labels = result_model(batch, interp)
pred1_target, pred2_target, labels_target = result_model(
batch_target, interp_target)
pred1_target_D = F.softmax((pred1_target), dim=1)
pred2_target_D = F.softmax((pred2_target), dim=1)
pred1_D = F.softmax((pred1), dim=1)
pred2_D = F.softmax((pred2), dim=1)
mix1_D = torch.cat((pred1_target_D, pred1_D), dim=1)
mix2_D = torch.cat((pred2_target_D, pred2_D), dim=1)
D_out_mix_1 = model_D1(mix1_D)
D_out_mix_2 = model_D2(mix2_D)
# D_out1 = D_out1.cuda(gpu_id_1)
# D_out2 = D_out2.cuda(gpu_id_1)
loss_adv_mix1 = bce_loss(
D_out_mix_1,
Variable(
torch.FloatTensor(D_out_mix_1.data.size()).fill_(
source_label)).cuda(gpu_id_1))
loss_adv_mix2 = bce_loss(
D_out_mix_2,
Variable(
torch.FloatTensor(D_out_mix_2.data.size()).fill_(
source_label)).cuda(gpu_id_1))
loss_adv_target1 = loss_adv_mix1 * 2
loss_adv_target2 = loss_adv_mix2 * 2
loss = args.lambda_adv_target1 * loss_adv_target1.cuda(
gpu_id_1) + args.lambda_adv_target2 * loss_adv_target2.cuda(
gpu_id_1)
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
pred1_last, pred2_last, labels_last = result_model(
batch_last, interp)
# train with source
pred1 = pred1.detach().cuda(gpu_id_1)
pred2 = pred2.detach().cuda(gpu_id_1)
pred1_target = pred1_target.detach().cuda(gpu_id_1)
pred2_target = pred2_target.detach().cuda(gpu_id_1)
pred1_last = pred1_last.detach().cuda(gpu_id_1)
pred2_last = pred2_last.detach().cuda(gpu_id_1)
pred1_D = F.softmax((pred1), dim=1)
pred2_D = F.softmax((pred2), dim=1)
pred1_last_D = F.softmax((pred1_last), dim=1)
pred2_last_D = F.softmax((pred2_last), dim=1)
pred1_target_D = F.softmax((pred1_target), dim=1)
pred2_target_D = F.softmax((pred2_target), dim=1)
real1_D = torch.cat((pred1_D, pred1_last_D), dim=1)
real2_D = torch.cat((pred2_D, pred2_last_D), dim=1)
mix1_D_ = torch.cat((pred1_last_D, pred1_target_D), dim=1)
mix2_D_ = torch.cat((pred2_last_D, pred2_target_D), dim=1)
D_out1_real = model_D1(real1_D)
D_out2_real = model_D2(real2_D)
D_out1_mix = model_D1(mix1_D_)
D_out2_mix = model_D2(mix2_D_)
# D_out1 = D_out1.cuda(gpu_id_1)
# D_out2 = D_out2.cuda(gpu_id_1)
loss_D1 = bce_loss(
D_out1_real,
Variable(
torch.FloatTensor(D_out1_real.data.size()).fill_(
source_label)).cuda(gpu_id_1))
loss_D2 = bce_loss(
D_out2_real,
Variable(
torch.FloatTensor(D_out2_real.data.size()).fill_(
source_label)).cuda(gpu_id_1))
loss_D3 = bce_loss(
D_out1_mix,
Variable(
torch.FloatTensor(D_out1_mix.data.size()).fill_(
mix_label)).cuda(gpu_id_1))
loss_D4 = bce_loss(
D_out2_mix,
Variable(
torch.FloatTensor(D_out2_mix.data.size()).fill_(
mix_label)).cuda(gpu_id_1))
loss_D1 = (loss_D1 + loss_D3) / args.iter_size / 2
loss_D2 = (loss_D2 + loss_D4) / 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
pred1 = pred1.detach().cuda(gpu_id_1)
pred2 = pred2.detach().cuda(gpu_id_1)
pred1_target = pred1_target.detach().cuda(gpu_id_1)
pred2_target = pred2_target.detach().cuda(gpu_id_1)
pred1_last_target = pred1_last_target.detach().cuda(gpu_id_1)
pred2_last_target = pred2_last_target.detach().cuda(gpu_id_1)
pred1_D = F.softmax((pred1), dim=1)
pred2_D = F.softmax((pred2), dim=1)
pred1_last_target_D = F.softmax((pred1_last_target), dim=1)
pred2_last_target_D = F.softmax((pred2_last_target), dim=1)
pred1_target_D = F.softmax((pred1_target), dim=1)
pred2_target_D = F.softmax((pred2_target), dim=1)
fake1_D_ = torch.cat((pred1_target_D, pred1_target_D), dim=1)
fake2_D_ = torch.cat((pred2_target_D, pred2_target_D), dim=1)
mix1_D__ = torch.cat((pred1_D, pred1_last_target_D), dim=1)
mix2_D__ = torch.cat((pred2_D, pred2_last_target_D), dim=1)
# pred_target1 = pred_target1.detach().cuda(gpu_id_1)
# pred_target2 = pred_target2.detach().cuda(gpu_id_1)
D_out1 = model_D1(fake1_D_)
D_out2 = model_D2(fake2_D_)
D_out3 = model_D1(mix1_D__)
D_out4 = model_D2(mix2_D__)
# D_out1 = D_out1.cuda(gpu_id_1)
# D_out2 = D_out2.cuda(gpu_id_1)
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(D_out1.data.size()).fill_(
target_label)).cuda(gpu_id_1))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(D_out2.data.size()).fill_(
target_label)).cuda(gpu_id_1))
loss_D3 = bce_loss(
D_out3,
Variable(
torch.FloatTensor(
D_out3.data.size()).fill_(mix_label)).cuda(gpu_id_1))
loss_D4 = bce_loss(
D_out4,
Variable(
torch.FloatTensor(
D_out4.data.size()).fill_(mix_label)).cuda(gpu_id_1))
loss_D1 = (loss_D1 + loss_D3) / args.iter_size / 2
loss_D2 = (loss_D2 + loss_D4) / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
batch_last, batch_last_target = batch, batch_target
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_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:
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."""
device = torch.device("cuda" if not args.cpu else "cpu")
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)
cudnn.enabled = True
# Create network
if args.model == 'DeepLab':
model = DeeplabMulti(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.to(device)
cudnn.benchmark = True
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(cityscapesDataSet(
max_iters=args.num_steps * args.iter_size * args.batch_size,
scale=False),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
seg_loss = torch.nn.CrossEntropyLoss(ignore_index=255)
interp = nn.Upsample(size=(input_size_target[1], input_size_target[0]),
mode='bilinear',
align_corners=True)
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)
_, batch = trainloader_iter.__next__()
images, labels, _, _ = batch
images = images.to(device)
labels = labels.long().to(device)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
loss_seg1 = seg_loss(pred1, labels)
loss_seg2 = seg_loss(pred2, labels)
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 source
pred1 = pred1.detach()
pred2 = pred2.detach()
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_stop) + '.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'))
def main():
"""Create the model and start the training."""
if RESTART:
args.snapshot_dir = RESTART_FROM
else:
args.snapshot_dir = generate_snapshot_name(args)
args_dict = vars(args)
import json
###### load args for restart ######
if RESTART:
# pdb.set_trace()
args_dict_file = args.snapshot_dir + 'args_dict_{}.json'.format(
RESTART_ITER)
with open(args_dict_file) as f:
args_dict_last = json.load(f)
for arg in args_dict:
args_dict[arg] = args_dict_last[arg]
###### load args for restart ######
device = torch.device("cuda" if not args.cpu else "cpu")
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)
cudnn.enabled = True
cudnn.benchmark = True
if args.model == 'DeepLab':
model = DeeplabMulti(num_classes=args.num_classes)
model_D1 = FCDiscriminator(num_classes=args.num_classes).to(device)
model_D2 = FCDiscriminator(num_classes=args.num_classes).to(device)
#### restore model_D1, D2 and model
if RESTART:
# pdb.set_trace()
# model parameters
restart_from_model = args.restart_from + 'GTA5_{}.pth'.format(
RESTART_ITER)
saved_state_dict = torch.load(restart_from_model)
model.load_state_dict(saved_state_dict)
# model_D1 parameters
restart_from_D1 = args.restart_from + 'GTA5_{}_D1.pth'.format(
RESTART_ITER)
saved_state_dict = torch.load(restart_from_D1)
model_D1.load_state_dict(saved_state_dict)
# model_D2 parameters
restart_from_D2 = args.restart_from + 'GTA5_{}_D2.pth'.format(
RESTART_ITER)
saved_state_dict = torch.load(restart_from_D2)
model_D2.load_state_dict(saved_state_dict)
#### model_D1, D2 are randomly initialized, model is pre-trained ResNet on ImageNet
else:
# model parameters
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.to(device)
model_D1.train()
model_D1.to(device)
model_D2.train()
model_D2.to(device)
#### From here, code should not be related to model reload ####
# but we would need hyperparameters: n_iter,
# [lr, momentum, weight_decay, betas](these are all in args)
# args.snapshot_dir = generate_snapshot_name()
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)
# pdb.set_trace()
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()
if args.gan == 'Vanilla':
bce_loss = torch.nn.BCEWithLogitsLoss()
elif args.gan == 'LS':
bce_loss = torch.nn.MSELoss()
seg_loss = torch.nn.CrossEntropyLoss(ignore_index=255)
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
# set up tensor board
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
for i_iter in range(args.start_steps, 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 = images.to(device)
labels = labels.long().to(device)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
pdb.set_trace()
loss_seg1 = seg_loss(pred1, labels)
loss_seg2 = seg_loss(pred2, labels)
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 = targetloader_iter.__next__()
images, _, _ = batch
images = images.to(device)
pdb.set_trace()
pred_target1, pred_target2 = model(images)
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
pdb.set_trace()
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_adv_target1 = bce_loss(
D_out1,
torch.FloatTensor(
D_out1.data.size()).fill_(source_label).to(device))
loss_adv_target2 = bce_loss(
D_out2,
torch.FloatTensor(
D_out2.data.size()).fill_(source_label).to(device))
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.item() / args.iter_size
loss_adv_target_value2 += loss_adv_target2.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,
torch.FloatTensor(
D_out1.data.size()).fill_(source_label).to(device))
loss_D2 = bce_loss(
D_out2,
torch.FloatTensor(
D_out2.data.size()).fill_(source_label).to(device))
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.item()
loss_D_value2 += loss_D2.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,
torch.FloatTensor(
D_out1.data.size()).fill_(target_label).to(device))
loss_D2 = bce_loss(
D_out2,
torch.FloatTensor(
D_out2.data.size()).fill_(target_label).to(device))
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.item()
loss_D_value2 += loss_D2.item()
optimizer.step()
optimizer_D1.step()
optimizer_D2.step()
scalar_info = {
'loss_seg1': loss_seg_value1,
'loss_seg2': loss_seg_value2,
'loss_adv_target1': loss_adv_target_value1,
'loss_adv_target2': loss_adv_target_value2,
'loss_D1': loss_D_value1,
'loss_D2': loss_D_value2,
}
if i_iter % 10 == 0:
for key, val in scalar_info.items():
writer.add_scalar(key, val, i_iter)
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_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'))
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'))
###### also record latest saved iteration #######
args_dict['learning_rate'] = optimizer.param_groups[0]['lr']
args_dict['learning_rate_D'] = optimizer_D1.param_groups[0]['lr']
args_dict['start_steps'] = i_iter
args_dict_file = args.snapshot_dir + '/args_dict_{}.json'.format(
i_iter)
with open(args_dict_file, 'w') as f:
json.dump(args_dict, f)
###### also record latest saved iteration #######
writer.close()
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 = DeeplabMulti(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_D = FCDiscriminator(num_classes=args.num_classes)
model_D.train()
model_D.cuda(0)
#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)
gta_trainloader = data.DataLoader(GTA5DataSet(
args.data_dir,
args.data_list,
args.num_classes,
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)
gta_trainloader_iter = enumerate(gta_trainloader)
gta_valloader = data.DataLoader(GTA5DataSet(args.data_dir,
args.valdata_list,
args.num_classes,
max_iters=None,
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)
gta_valloader_iter = enumerate(gta_valloader)
cityscapes_targetloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target,
args.num_classes,
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='train'),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
cityscapes_targetloader_iter = enumerate(cityscapes_targetloader)
cityscapes_valtargetloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.valdata_list_target,
args.num_classes,
max_iters=None,
crop_size=input_size_target,
scale=False,
mirror=args.random_mirror,
mean=IMG_MEAN,
set='val'),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
cityscapes_valtargetloader_iter = enumerate(cityscapes_valtargetloader)
# 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_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.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
for i_iter in range(args.num_steps):
loss_D_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_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
_, batch = gta_trainloader_iter.next()
images, labels, _, _ = batch
size = labels.size()
#print(size)
#labels = Variable(labels)
oneHot_size = (size[0], args.num_classes, size[2], size[3])
input_label = torch.FloatTensor(torch.Size(oneHot_size)).zero_()
input_label = input_label.scatter_(1, labels.long(), 1.0)
#print(input_label.size())
labels1 = Variable(input_label).cuda(0)
#D_out1 = model_D(labels)
#print(D_out1.data.size())
#loss_out1 = bce_loss(D_out1, Variable(torch.FloatTensor(D_out1.data.size()).fill_(source_label)).cuda(0))
_, batch = cityscapes_targetloader_iter.next()
images, labels, _, _ = batch
size = labels.size()
#labels = Variable(labels)
oneHot_size = (size[0], args.num_classes, size[2], size[3])
input_label = torch.FloatTensor(torch.Size(oneHot_size)).zero_()
input_label = input_label.scatter_(1, labels.long(), 1.0)
labels2 = Variable(input_label).cuda(0)
#print(labels1.data.size())
#print(labels2.data.size())
labels = torch.cat((labels1, labels2), 0)
#print(labels.data.size())
#D_out2 = model_D(labels)
D_out = model_D(labels)
#print(D_out.data.size())
target_size = D_out.data.size()
target_labels1 = torch.FloatTensor(
torch.Size((target_size[0] / 2, target_size[1], target_size[2],
target_size[3]))).fill_(source_label)
target_labels2 = torch.FloatTensor(
torch.Size((target_size[0] / 2, target_size[1], target_size[2],
target_size[3]))).fill_(target_label)
target_labels = torch.cat((target_labels1, target_labels2), 0)
target_labels = Variable(target_labels).cuda(0)
#print(target_labels.data.size())
loss_out = bce_loss(D_out, target_labels)
loss = loss_out / args.iter_size
loss.backward()
loss_D_value += loss_out.data.cpu().numpy()[0] / args.iter_size
#print(loss_D_value)
optimizer_D.step()
#print('exp = {}'.format(args.snapshot_dir))
if i_iter % 100 == 0:
print('iter = {0:8d}/{1:8d}, loss_D = {2:.3f}'.format(
i_iter, args.num_steps, loss_D_value))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'Classify_' + str(args.num_steps) + '.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'Classify_' + str(i_iter) + '.pth'))
model_D.eval()
loss_valD_value = 0
correct = 0
wrong = 0
for i, (images, labels, _, _) in enumerate(gta_valloader):
#if i > 500:
# break
size = labels.size()
#labels = Variable(labels)
oneHot_size = (size[0], args.num_classes, size[2], size[3])
input_label = torch.FloatTensor(
torch.Size(oneHot_size)).zero_()
input_label = input_label.scatter_(1, labels.long(), 1.0)
labels = Variable(input_label).cuda(0)
D_out1 = model_D(labels)
loss_out1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(0))
loss_valD_value += loss_out1.data.cpu().numpy()[0]
correct = correct + (D_out1.data.cpu() < 0).sum() / 100
wrong = wrong + (D_out1.data.cpu() >= 0).sum() / 100
#accuracy = 1.0 * correct / (wrong + correct)
#print('accuracy:%f' % accuracy)
#print(correct)
#print(wrong)
for i, (images, labels, _,
_) in enumerate(cityscapes_valtargetloader):
#if i > 500:
# break
size = labels.size()
#labels = Variable(labels)
oneHot_size = (size[0], args.num_classes, size[2], size[3])
input_label = torch.FloatTensor(
torch.Size(oneHot_size)).zero_()
input_label = input_label.scatter_(1, labels.long(), 1.0)
labels = Variable(input_label).cuda(0)
D_out2 = model_D(labels)
loss_out2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(target_label)).cuda(0))
loss_valD_value += loss_out2.data.cpu().numpy()[0]
wrong = wrong + (D_out2.data.cpu() < 0).sum()
correct = correct + (D_out2.data.cpu() >= 0).sum()
accuracy = 1.0 * correct / (wrong + correct)
print('accuracy:%f' % accuracy)
print('iter = {0:8d}/{1:8d}, loss_valD = {2:.3f}'.format(
i_iter, args.num_steps, loss_valD_value))
model_D.train()
def main():
"""Create the model and start the training."""
h, w = map(int, args.input_size_source.split(','))
input_size_source = (h, w)
h, w = map(int, args.input_size_target.split(','))
input_size_target = (h, w)
cudnn.enabled = True
# Create Network
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:
i_parts = i.split('.')
if not args.num_classes == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
if args.restore_from[:4] == './mo':
model.load_state_dict(new_params)
else:
model.load_state_dict(saved_state_dict)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# Init D
model_D = FCDiscriminator(num_classes=args.num_classes)
# =============================================================================
# #for retrain
# saved_state_dict_D = torch.load(RESTORE_FROM_D)
# model_D.load_state_dict(saved_state_dict_D)
# =============================================================================
model_D.train()
model_D.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
if args.source == '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_source,
scale=True,
mirror=True,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
else:
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_source,
scale=True,
mirror=True,
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=True,
mirror=True,
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)
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
weighted_bce_loss = WeightedBCEWithLogitsLoss()
interp_source = nn.Upsample(size=(input_size_source[1],
input_size_source[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):
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
damping = (1 - i_iter / NUM_STEPS)
#======================================================================================
# train G
#======================================================================================
#Remove Grads in D
for param in model_D.parameters():
param.requires_grad = False
# Train with Source
_, batch = next(trainloader_iter)
images_s, labels_s, _, _, _ = batch
images_s = Variable(images_s).cuda(args.gpu)
pred_source1, pred_source2 = model(images_s)
pred_source1 = interp_source(pred_source1)
pred_source2 = interp_source(pred_source2)
#Segmentation Loss
loss_seg = (loss_calc(pred_source1, labels_s, args.gpu) +
loss_calc(pred_source2, labels_s, args.gpu))
loss_seg.backward()
# Train with Target
_, batch = next(targetloader_iter)
images_t, _, _, _ = batch
images_t = Variable(images_t).cuda(args.gpu)
pred_target1, pred_target2 = model(images_t)
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
weight_map = weightmap(F.softmax(pred_target1, dim=1),
F.softmax(pred_target2, dim=1))
D_out = interp_target(
model_D(F.softmax(pred_target1 + pred_target2, dim=1)))
#Adaptive Adversarial Loss
if (i_iter > PREHEAT_STEPS):
loss_adv = weighted_bce_loss(
D_out,
Variable(
torch.FloatTensor(
D_out.data.size()).fill_(source_label)).cuda(args.gpu),
weight_map, Epsilon, Lambda_local)
else:
loss_adv = bce_loss(
D_out,
Variable(
torch.FloatTensor(
D_out.data.size()).fill_(source_label)).cuda(args.gpu))
loss_adv = loss_adv * Lambda_adv * damping
loss_adv.backward()
#Weight Discrepancy Loss
W5 = None
W6 = None
if args.model == 'ResNet':
for (w5, w6) in zip(model.layer5.parameters(),
model.layer6.parameters()):
if W5 is None and W6 is None:
W5 = w5.view(-1)
W6 = w6.view(-1)
else:
W5 = torch.cat((W5, w5.view(-1)), 0)
W6 = torch.cat((W6, w6.view(-1)), 0)
loss_weight = (torch.matmul(W5, W6) /
(torch.norm(W5) * torch.norm(W6)) + 1
) # +1 is for a positive loss
loss_weight = loss_weight * Lambda_weight * damping * 2
loss_weight.backward()
#======================================================================================
# train D
#======================================================================================
# Bring back Grads in D
for param in model_D.parameters():
param.requires_grad = True
# Train with Source
pred_source1 = pred_source1.detach()
pred_source2 = pred_source2.detach()
D_out_s = interp_source(
model_D(F.softmax(pred_source1 + pred_source2, dim=1)))
loss_D_s = bce_loss(
D_out_s,
Variable(
torch.FloatTensor(
D_out_s.data.size()).fill_(source_label)).cuda(args.gpu))
loss_D_s.backward()
# Train with Target
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
weight_map = weight_map.detach()
D_out_t = interp_target(
model_D(F.softmax(pred_target1 + pred_target2, dim=1)))
#Adaptive Adversarial Loss
if (i_iter > PREHEAT_STEPS):
loss_D_t = weighted_bce_loss(
D_out_t,
Variable(
torch.FloatTensor(
D_out_t.data.size()).fill_(target_label)).cuda(
args.gpu), weight_map, Epsilon, Lambda_local)
else:
loss_D_t = bce_loss(
D_out_t,
Variable(
torch.FloatTensor(
D_out_t.data.size()).fill_(target_label)).cuda(
args.gpu))
loss_D_t.backward()
optimizer.step()
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:6d}/{1:6d}, loss_seg = {2:.4f} loss_adv = {3:.4f}, loss_weight = {4:.4f}, loss_D_s = {5:.4f} loss_D_t = {6:.4f}'
.format(i_iter, args.num_steps, loss_seg, loss_adv, loss_weight,
loss_D_s, loss_D_t))
f_loss = open(osp.join(args.snapshot_dir, 'loss.txt'), 'a')
f_loss.write('{0:.4f} {1:.4f} {2:.4f} {3:.4f} {4:.4f}\n'.format(
loss_seg, loss_adv, loss_weight, loss_D_s, loss_D_t))
f_loss.close()
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_D.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps) + '_D.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_D.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D.pth'))
def main():
"""Create the model and start the evaluation process."""
import warnings
if not sys.warnoptions:
warnings.simplefilter("ignore")
args = get_arguments()
gpu0 = args.gpu
if not os.path.exists(args.save):
os.makedirs(args.save)
model = Res_Deeplab(num_classes=args.num_classes)
#from pytorchgo.model.MyFCN8s import MyFCN8s
#model = MyFCN8s(n_class=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_state_dict'])
model.eval()
model.cuda(gpu0)
image_size = (1024, 512)#(1280,720) #(2048, 1024)
cityscape_image_size = (2048, 1024)
print ("evaluating {}".format(args.restore_from))
print ("************ best mIoU:{} *******".format(saved_state_dict['best_mean_iu']))
print("evaluation image size: {}, please make sure this image size is equal to your training image size, this is important for your final mIoU!".format(image_size))
testloader = data.DataLoader(cityscapesDataSet( crop_size=(image_size[0], image_size[1]), mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
batch_size=1, shuffle=False, pin_memory=True)
interp = nn.Upsample(size=(cityscape_image_size[1], cityscape_image_size[0]), mode='bilinear')
from tensorpack.utils.stats import MIoUStatistics
stat = MIoUStatistics(NUM_CLASSES)
for index, batch in tqdm(enumerate(testloader)):
image,label, _, name = batch
image, label = Variable(image, volatile=True), Variable(label)
#output2 = model(image.cuda(gpu0))
output1, output2 = model(image.cuda(gpu0))
output = interp(output2).cpu().data[0].numpy()
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
name = name[0].split('/')[-1]
output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
stat.feed(output, label.data.cpu().numpy().squeeze())
print("tensorpack IoU: {}".format(stat.mIoU_beautify))
print("tensorpack class16 IoU: {}".format(np.sum(stat.IoU)/16))
print("tensorpack mIoU: {}".format(stat.mIoU))
print("tensorpack mean_accuracy: {}".format(stat.mean_accuracy))
print("tensorpack accuracy: {}".format(stat.accuracy))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
config_path = os.path.join(os.path.dirname(args.restore_from), 'opts.yaml')
with open(config_path, 'r') as stream:
config = yaml.load(stream)
args.model = config['model']
print('ModelType:%s' % args.model)
print('NormType:%s' % config['norm_style'])
gpu0 = args.gpu
batchsize = args.batchsize
model_name = os.path.basename(os.path.dirname(args.restore_from))
args.save += model_name
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeepLab':
model = DeeplabMulti(num_classes=args.num_classes,
use_se=config['use_se'],
train_bn=False,
norm_style=config['norm_style'])
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
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)
else:
saved_state_dict = torch.load(args.restore_from)
try:
model.load_state_dict(saved_state_dict)
except:
model = torch.nn.DataParallel(model)
model.load_state_dict(saved_state_dict)
#model = torch.nn.DataParallel(model)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(cityscapesDataSet(args.data_dir,
args.data_list,
crop_size=(512, 1024),
resize_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=batchsize,
shuffle=False,
pin_memory=True,
num_workers=4)
scale = 1.25
testloader2 = data.DataLoader(cityscapesDataSet(
args.data_dir,
args.data_list,
crop_size=(round(512 * scale), round(1024 * scale)),
resize_size=(round(1024 * scale), round(512 * scale)),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=batchsize,
shuffle=False,
pin_memory=True,
num_workers=4)
scale = 0.9
testloader3 = data.DataLoader(cityscapesDataSet(
args.data_dir,
args.data_list,
crop_size=(round(512 * scale), round(1024 * scale)),
resize_size=(round(1024 * scale), round(512 * scale)),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=batchsize,
shuffle=False,
pin_memory=True,
num_workers=4)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
sm = torch.nn.Softmax(dim=1)
log_sm = torch.nn.LogSoftmax(dim=1)
kl_distance = nn.KLDivLoss(reduction='none')
for index, img_data in enumerate(zip(testloader, testloader2,
testloader3)):
batch, batch2, batch3 = img_data
image, _, _, name = batch
image2, _, _, name2 = batch2
#image3, _, _, name3 = batch3
inputs = image.cuda()
inputs2 = image2.cuda()
#inputs3 = Variable(image3).cuda()
print('\r>>>>Extracting feature...%03d/%03d' %
(index * batchsize, NUM_STEPS),
end='')
if args.model == 'DeepLab':
with torch.no_grad():
output1, output2 = model(inputs)
output_batch = interp(sm(0.5 * output1 + output2))
heatmap_output1, heatmap_output2 = output1, output2
#output_batch = interp(sm(output1))
#output_batch = interp(sm(output2))
output1, output2 = model(fliplr(inputs))
output1, output2 = fliplr(output1), fliplr(output2)
output_batch += interp(sm(0.5 * output1 + output2))
heatmap_output1, heatmap_output2 = heatmap_output1 + output1, heatmap_output2 + output2
#output_batch += interp(sm(output1))
#output_batch += interp(sm(output2))
del output1, output2, inputs
output1, output2 = model(inputs2)
output_batch += interp(sm(0.5 * output1 + output2))
#output_batch += interp(sm(output1))
#output_batch += interp(sm(output2))
output1, output2 = model(fliplr(inputs2))
output1, output2 = fliplr(output1), fliplr(output2)
output_batch += interp(sm(0.5 * output1 + output2))
#output_batch += interp(sm(output1))
#output_batch += interp(sm(output2))
del output1, output2, inputs2
output_batch = output_batch.cpu().data.numpy()
heatmap_batch = torch.sum(kl_distance(log_sm(heatmap_output1),
sm(heatmap_output2)),
dim=1)
heatmap_batch = torch.log(
1 + 10 * heatmap_batch) # for visualization
heatmap_batch = heatmap_batch.cpu().data.numpy()
#output1, output2 = model(inputs3)
#output_batch += interp(sm(0.5* output1 + output2)).cpu().data.numpy()
#output1, output2 = model(fliplr(inputs3))
#output1, output2 = fliplr(output1), fliplr(output2)
#output_batch += interp(sm(0.5 * output1 + output2)).cpu().data.numpy()
#del output1, output2, inputs3
elif args.model == 'DeeplabVGG' or args.model == 'Oracle':
output_batch = model(Variable(image).cuda())
output_batch = interp(output_batch).cpu().data.numpy()
output_batch = output_batch.transpose(0, 2, 3, 1)
scoremap_batch = np.asarray(np.max(output_batch, axis=3))
output_batch = np.asarray(np.argmax(output_batch, axis=3),
dtype=np.uint8)
output_iterator = []
heatmap_iterator = []
scoremap_iterator = []
for i in range(output_batch.shape[0]):
output_iterator.append(output_batch[i, :, :])
heatmap_iterator.append(heatmap_batch[i, :, :] /
np.max(heatmap_batch[i, :, :]))
scoremap_iterator.append(1 - scoremap_batch[i, :, :] /
np.max(scoremap_batch[i, :, :]))
name_tmp = name[i].split('/')[-1]
name[i] = '%s/%s' % (args.save, name_tmp)
with Pool(4) as p:
p.map(save, zip(output_iterator, name))
p.map(save_heatmap, zip(heatmap_iterator, name))
p.map(save_scoremap, zip(scoremap_iterator, name))
del output_batch
return args.save
