def check_original_discriminator(args, pred_target1, pred_target2, n_iter):
device = torch.device("cuda" if not args.cpu else "cpu")
model_D1 = FCDiscriminator(num_classes=args.num_classes).to(device)
model_D2 = FCDiscriminator(num_classes=args.num_classes).to(device)
model_D1.eval()
model_D2.eval()
baseline_dir = 'snapshots/baseline_single_250000_seg0.1_adv10.0002_adv20.001_bs1_11-10-8-52/'
bce_loss = torch.nn.BCEWithLogitsLoss()
#### restore model_D1, D2 and model
# model_D1 parameters
D1 = baseline_dir + 'GTA5_150000_D1.pth'
saved_state_dict = torch.load(D1)
model_D1.load_state_dict(saved_state_dict)
# model_D2 parameters
D2 = baseline_dir + 'GTA5_150000_D2.pth'
saved_state_dict = torch.load(D2)
model_D2.load_state_dict(saved_state_dict)
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
orig_d1_loss = bce_loss(
D_out1,
torch.FloatTensor(D_out1.data.size()).fill_(1).to(device))
orig_d2_loss = bce_loss(
D_out2,
torch.FloatTensor(D_out2.data.size()).fill_(1).to(device))
original_discriminator_file = args.snapshot_dir + '/original_discriminator.txt'
with open(original_discriminator_file, 'a+') as f:
f.write('n_iter:{} '.format(n_iter))
f.write('orig d1 loss:{} '.format(orig_d1_loss))
f.write('orig d2 loss:{} '.format(orig_d2_loss))
f.write('\n')
del model_D1
del model_D2
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 evaluation process."""
args = get_arguments()
if not os.path.exists(args.save):
os.makedirs(args.save)
gpu0 = args.gpu
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)
#===========load discriminator model====begin===
model_d2 = FCDiscriminator(num_classes=args.num_classes)
d2_state_dict = torch.load(args.dis_restore_from)
model_d2.load_state_dict(d2_state_dict)
model_d2.eval()
model_d2.cuda(gpu0)
#===========load discriminator model====end===
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)
# interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear')
out_values = []
fine_out_values = []
retrain_list = []
file = open(CITYS_RETRAIN_TXT, 'w')
for index, batch in enumerate(testloader):
if index % 20 == 0:
print('%d processd of %d' % (index, len(testloader)))
image, _, name = batch
output1, output2 = model(Variable(image, volatile=True).cuda(gpu0))
ini_output = interp(output2)
d2_out1 = model_d2(F.softmax(ini_output,
dim=1)) #.cpu().data[0].numpy()
out_valu = d2_out1.mean()
out_valu_img = np.array([[name[0]], out_valu.cpu().data.numpy()])
out_values.extend(out_valu_img)
if out_valu.cpu().data.numpy() > 0.64:
fine_out_valu_img = np.array([[name[0]],
out_valu.cpu().data.numpy()])
fine_out_values.extend(fine_out_valu_img)
file.write(name[0] + '\n')
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]))
output_col.save('%s/%s.png' % (args.save, name.split('.')[0]))
# print('its confidence value is %f' % out_valu)
# plt.imshow(output_col)
# plt.title(str(out_valu))
# plt.show()
# output = Image.fromarray(output)
# output.save('%s/%s' % (args.save, name))
out_values = np.array(out_values)
np.save(CITYS_VALUES_SV_PATH, out_values)
np.save(CITYS_FINE_VALUES_SV_PATH, fine_out_values)
file.close()
