def main():
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
kernel_size = (321, 321)
cudnn.enabled = True
gpu = args.gpu
np.random.seed(args.random_seed)
# create network
model = Res_Deeplab(num_classes=args.num_classes)
# load pretrained 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)
# only copy the params that exist in current model (caffe-like)
new_params = model.state_dict().copy()
for name, param in new_params.items():
print(name)
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
print('copy {}'.format(name))
model.load_state_dict(new_params)
model.train()
model.cuda(args.gpu)
# init D
model_D = FCDiscriminator(num_classes=args.num_classes,
kernel_size=kernel_size)
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D.train()
model_D.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
# load dataset
train_dataset = VOCDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
if args.partial_data is None:
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
else:
#sample partial data
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print('loading train ids from {}'.format(args.partial_id))
else:
train_ids = np.arange(train_dataset_size)
np.random.shuffle(train_ids)
pickle.dump(train_ids,
open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb'))
# labeled data
train_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
train_gt_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=3,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
sampler=train_gt_sampler,
num_workers=3,
pin_memory=True)
# unlabeled data
train_remain_sampler = data.sampler.SubsetRandomSampler(
train_ids[partial_size:])
trainloader_remain = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_remain_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
# loss/bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear')
# labels for adversarial training
pred_label = 0
gt_label = 1
for i_iter in range(args.num_steps):
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_adv_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 Segmentation
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# do semi-supervised training first
if args.lambda_semi_adv > 0 and i_iter >= args.semi_start_adv:
try:
_, batch = trainloader_remain_iter.next()
except:
trainloader_remain_iter = enumerate(trainloader_remain)
_, batch = trainloader_remain_iter.next()
# only access to img
images, _, _, _ = batch
images = Variable(images).cuda(args.gpu)
# prediction and pooling
pred = interp(model(images))
pred_remain = pred.detach()
D_out = interp(model_D(F.softmax(pred)))
D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(
axis=1)
ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(
np.bool)
# adversarial loss
loss_semi_adv = args.lambda_semi_adv * bce_loss(
D_out, make_D_label(gt_label, ignore_mask_remain))
loss_semi_adv = loss_semi_adv / args.iter_size
# true loss value without multiplier
loss_semi_adv_value += loss_semi_adv.data.cpu().numpy(
) / args.lambda_semi_adv
loss_semi_adv.backward()
else:
loss_semi_adv = None
# train with labeled images
try:
_, batch = trainloader_iter.next()
except:
trainloader_iter = enumerate(trainloader)
_, batch = trainloader_iter.next()
images, labels, _, _ = batch
images = Variable(images).cuda(args.gpu)
ignore_mask = (labels.numpy() == 255)
# prediction + pooling
pred = interp(model(images))
D_out = interp(model_D(F.softmax(pred)))
# computing loss
loss_seg = loss_calc(pred, labels, args.gpu)
loss_adv_pred = bce_loss(D_out,
make_D_label(gt_label, ignore_mask))
loss = loss_seg + args.lambda_adv_pred * loss_adv_pred
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size
loss_adv_pred_value += loss_adv_pred.data.cpu().numpy(
) / args.iter_size
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
if args.D_remain and args.lambda_semi_adv > 0:
pred = torch.cat((pred, pred_remain), 0)
ignore_mask = np.concatenate((ignore_mask, ignore_mask_remain),
axis=0)
D_out = interp(model_D(F.softmax(pred)))
loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()
# train with gt
# get gt labels
try:
_, batch = trainloader_gt_iter.next()
except:
trainloader_gt_iter = enumerate(trainloader_gt)
_, batch = trainloader_gt_iter.next()
_, labels_gt, _, _ = batch
D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu)
ignore_mask_gt = (labels_gt.numpy() == 255)
D_out = interp(model_D(D_gt_v))
loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()
optimizer.step()
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_semi_adv = {5:.3f}'
.format(i_iter, args.num_steps, loss_seg_value,
loss_adv_pred_value, loss_D_value, loss_semi_adv_value))
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(
args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_' +
str(args.random_seed) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(
args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_' +
str(args.random_seed) + '_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, 'VOC_' + str(i_iter) + '_' +
str(args.random_seed) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(
args.snapshot_dir, 'VOC_' + str(i_iter) + '_' +
str(args.random_seed) + '_D.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
# 将参数的input_size 映射到整数，并赋值，从字符串转换到整数二元组
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = False
gpu = args.gpu
# create network
model = Res_Deeplab(num_classes=args.num_classes)
# load pretrained 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)
# only copy the params that exist in current model (caffe-like)
# 确保模型中参数的格式与要加载的参数相同
# 返回一个字典，保存着module的所有状态（state）；parameters和persistent buffers都会包含在字典中，字典的key就是parameter和buffer的 names。
new_params = model.state_dict().copy()
for name, param in new_params.items():
# print (name)
if name in saved_state_dict and param.size() == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
# print('copy {}'.format(name))
model.load_state_dict(new_params)
# 设置为训练模式
model.train()
cudnn.benchmark = True
model.cuda(gpu)
# init D
model_D = FCDiscriminator(num_classes=args.num_classes)
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D.train()
model_D.cuda(gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size,
scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size,
scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN)
if args.partial_data is None:
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True)
else:
# sample partial data
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print('loading train ids from {}'.format(args.partial_id))
else:
train_ids = list(range(train_dataset_size)) # ?
np.random.shuffle(train_ids)
pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) # 写入文件
train_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size])
train_remain_sampler = data.sampler.SubsetRandomSampler(train_ids[partial_size:])
train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size])
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True)
trainloader_remain = data.DataLoader(train_dataset,
batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3,
pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
optimizer_D.zero_grad()
# loss/ bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') # ???
# labels for adversarial training
pred_label = 0
gt_label = 1
for i_iter in range(args.num_steps):
print("Iter:", i_iter)
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_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
# do semi first
if args.lambda_semi > 0 and i_iter >= args.semi_start:
try:
_, batch = next(trainloader_remain_iter)
except:
trainloader_remain_iter = enumerate(trainloader_remain)
_, batch = next(trainloader_remain_iter)
# only access to img
images, _, _, _ = batch
images = Variable(images).cuda(gpu)
# images = Variable(images).cpu()
pred = interp(model(images))
D_out = interp(model_D(F.softmax(pred)))
D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(axis=1)
# produce ignore mask
semi_ignore_mask = (D_out_sigmoid < args.mask_T)
semi_gt = pred.data.cpu().numpy().argmax(axis=1)
semi_gt[semi_ignore_mask] = 255
semi_ratio = 1.0 - float(semi_ignore_mask.sum()) / semi_ignore_mask.size
print('semi ratio: {:.4f}'.format(semi_ratio))
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
loss_semi = args.lambda_semi * loss_calc(pred, semi_gt, args.gpu)
loss_semi = loss_semi / args.iter_size
loss_semi.backward()
loss_semi_value += loss_semi.data.cpu().numpy()[0] / args.lambda_semi
else:
loss_semi = None
# train with source
try:
_, batch = next(trainloader_iter)
except:
trainloader_iter = enumerate(trainloader)
_, batch = next(trainloader_iter)
images, labels, _, _ = batch
images = Variable(images).cuda(gpu)
# images = Variable(images).cpu()
ignore_mask = (labels.numpy() == 255)
pred = interp(model(images))
loss_seg = loss_calc(pred, labels, args.gpu)
D_out = interp(model_D(F.softmax(pred)))
loss_adv_pred = bce_loss(D_out, make_D_label(gt_label, ignore_mask))
loss = loss_seg + args.lambda_adv_pred * loss_adv_pred
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value += loss_seg.data.cpu().numpy()[0] / args.iter_size
loss_adv_pred_value += loss_adv_pred.data.cpu().numpy()[0] / args.iter_size
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
D_out = interp(model_D(F.softmax(pred)))
loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()[0]
# train with gt
# get gt labels
try:
_, batch = next(trainloader_gt_iter)
except:
trainloader_gt_iter = enumerate(trainloader_gt)
_, batch = next(trainloader_gt_iter)
_, labels_gt, _, _ = batch
D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu)
# D_gt_v = Variable(one_hot(labels_gt)).cpu()
ignore_mask_gt = (labels_gt.numpy() == 255)
D_out = interp(model_D(D_gt_v))
loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()[0]
optimizer.step()
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_semi = {5:.3f}'.format(
i_iter, args.num_steps, loss_seg_value, loss_adv_pred_value, loss_D_value, loss_semi_value))
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '.pth'))
torch.save(model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + 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, 'VOC_' + str(i_iter) + '.pth'))
torch.save(model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
gpu = args.gpu
# create network
model = Res_Deeplab(num_classes=args.num_classes)
# load pretrained 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)
# only copy the params that exist in current model (caffe-like)
new_params = model.state_dict().copy()
for name, param in new_params.items():
print(name)
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
print('copy {}'.format(name))
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)
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D.train()
model_D.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
train_dataset = VOCDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
if args.partial_data is None:
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=16,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=16,
pin_memory=True)
else:
#sample partial data
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print('loading train ids from {}'.format(args.partial_id))
else:
train_ids = np.arange(train_dataset_size)
np.random.shuffle(train_ids)
pickle.dump(train_ids,
open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb'))
train_sampler_all = data.sampler.SubsetRandomSampler(train_ids)
train_gt_sampler_all = data.sampler.SubsetRandomSampler(train_ids)
train_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
train_remain_sampler = data.sampler.SubsetRandomSampler(
train_ids[partial_size:])
train_gt_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
trainloader_all = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_sampler_all,
num_workers=16,
pin_memory=True)
trainloader_gt_all = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
sampler=train_gt_sampler_all,
num_workers=16,
pin_memory=True)
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=16,
pin_memory=True)
trainloader_remain = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_remain_sampler,
num_workers=16,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
sampler=train_gt_sampler,
num_workers=16,
pin_memory=True)
trainloader_remain_iter = iter(trainloader_remain)
trainloader_all_iter = iter(trainloader_all)
trainloader_iter = iter(trainloader)
trainloader_gt_iter = iter(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
# loss/ bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
# labels for adversarial training
pred_label = 0
gt_label = 1
#y_real_, y_fake_ = Variable(torch.ones(args.batch_size, 1).cuda()), Variable(torch.zeros(args.batch_size, 1).cuda())
for i_iter in range(args.num_steps):
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_fm_value = 0
loss_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
try:
batch = next(trainloader_iter)
except:
trainloader_iter = iter(trainloader)
batch = next(trainloader_iter)
images, labels, _, _ = batch
images = Variable(images).cuda(args.gpu)
#ignore_mask = (labels.numpy() == 255)
pred = interp(model(images))
loss_seg = loss_calc(pred, labels, args.gpu)
loss_seg.backward()
loss_seg_value += loss_seg.data.cpu().numpy()[0] / args.iter_size
if i_iter >= args.adv_start:
#fm loss calc
try:
batch = next(trainloader_all_iter)
except:
trainloader_iter = iter(trainloader_all)
batch = next(trainloader_all_iter)
images, labels, _, _ = batch
images = Variable(images).cuda(args.gpu)
#ignore_mask = (labels.numpy() == 255)
pred = interp(model(images))
_, D_out_y_pred = model_D(F.softmax(pred))
trainloader_gt_iter = iter(trainloader_gt)
batch = next(trainloader_gt_iter)
_, labels_gt, _, _ = batch
D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu)
#ignore_mask_gt = (labels_gt.numpy() == 255)
_, D_out_y_gt = model_D(D_gt_v)
fm_loss = torch.mean(
torch.abs(
torch.mean(D_out_y_gt, 0) -
torch.mean(D_out_y_pred, 0)))
loss = loss_seg + args.lambda_fm * fm_loss
# proper normalization
fm_loss.backward()
#loss_seg_value += loss_seg.data.cpu().numpy()[0]/args.iter_size
loss_fm_value += fm_loss.data.cpu().numpy()[0] / args.iter_size
loss_value += loss.data.cpu().numpy()[0] / args.iter_size
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
D_out_z, _ = model_D(F.softmax(pred))
y_fake_ = Variable(torch.zeros(D_out_z.size(0), 1).cuda())
loss_D_fake = criterion(D_out_z, y_fake_)
# train with gt
# get gt labels
_, labels_gt, _, _ = batch
D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu)
#ignore_mask_gt = (labels_gt.numpy() == 255)
D_out_z_gt, _ = model_D(D_gt_v)
#D_out = interp(D_out_x)
y_real_ = Variable(torch.ones(D_out_z_gt.size(0), 1).cuda())
loss_D_real = criterion(D_out_z_gt, y_real_)
loss_D = loss_D_fake + loss_D_real
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()[0]
optimizer.step()
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print('iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_D = {3:.3f}'.
format(i_iter, args.num_steps, loss_seg_value, loss_D_value))
print('fm_loss: ', loss_fm_value, ' g_loss: ', loss_value)
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + str(args.num_steps) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + 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, 'VOC_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
gpu = args.gpu
# create network
model = Res_Deeplab(num_classes=args.num_classes)
# load pretrained parameters (weights)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(
args.restore_from
) ## http://vllab1.ucmerced.edu/~whung/adv-semi-seg/resnet101COCO-41f33a49.pth
else:
saved_state_dict = torch.load(args.restore_from)
#checkpoint = torch.load(args.restore_from)_
# only copy the params that exist in current model (caffe-like)
new_params = model.state_dict().copy() # state_dict() is current model
for name, param in new_params.items():
#print (name) # 'conv1.weight, name:param(value), dict
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
#print('copy {}'.format(name))
model.load_state_dict(new_params)
#model.load_state_dict(checkpoint['state_dict'])
#optimizer.load_state_dict(args.checkpoint['optim_dict'])
model.train(
) # https://pytorch.org/docs/stable/nn.html, Sets the module in training mode.
model.cuda(args.gpu) ##
cudnn.benchmark = True # This flag allows you to enable the inbuilt cudnn auto-tuner to find the best algorithm to use for your hardware
# init D
model_D = FCDiscriminator(num_classes=args.num_classes)
#args.restore_from_D = 'snapshots/linear2/VOC_25000_D.pth'
if args.restore_from_D is not None: # None
model_D.load_state_dict(torch.load(args.restore_from_D))
# checkpoint_D = torch.load(args.restore_from_D)
# model_D.load_state_dict(checkpoint_D['state_dict'])
# optimizer_D.load_state_dict(checkpoint_D['optim_dict'])
model_D.train()
model_D.cuda(args.gpu)
if USECALI:
model_cali = ModelWithTemperature(model, model_D)
model_cali.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
random.seed(args.random_seed)
np.random.seed(args.random_seed)
torch.manual_seed(args.random_seed)
torch.cuda.manual_seed(args.random_seed)
train_dataset = VOCDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_dataset_remain = VOCDataSet(args.data_dir,
args.data_list_remain,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_dataset_size_remain = len(train_dataset_remain)
print train_dataset_size
print train_dataset_size_remain
train_gt_dataset = VOCGTDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
if args.partial_data is None: #if not partial, load all
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
else:
#sample partial data
#args.partial_data = 0.125
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print('loading train ids from {}'.format(args.partial_id))
else: #args.partial_id is none
train_ids = range(train_dataset_size)
train_ids_remain = range(train_dataset_size_remain)
np.random.shuffle(train_ids) #shuffle!
np.random.shuffle(train_ids_remain)
pickle.dump(train_ids,
open(osp.join(args.snapshot_dir, 'train_id.pkl'),
'wb')) #randomly suffled ids
#sampler
train_sampler = data.sampler.SubsetRandomSampler(
train_ids[:]) # 0~1/8,
train_remain_sampler = data.sampler.SubsetRandomSampler(
train_ids_remain[:])
train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:])
# train_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) # 0~1/8
# train_remain_sampler = data.sampler.SubsetRandomSampler(train_ids[partial_size:]) # used as unlabeled, 7/8
# train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size])
#train loader
trainloader = data.DataLoader(
train_dataset,
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=3,
pin_memory=True) # multi-process data loading
trainloader_remain = data.DataLoader(train_dataset_remain,
batch_size=args.batch_size,
sampler=train_remain_sampler,
num_workers=3,
pin_memory=True)
# trainloader_remain = data.DataLoader(train_dataset,
# batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3,
# pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
sampler=train_gt_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
# model.optim_paramters(args) = list(dict1, dict2), dict1 >> 'lr' and 'params'
# print(type(model.optim_parameters(args)[0]['params'])) # generator
#print(model.state_dict()['coeff'][0]) #confirmed
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
#optimizer.add_param_group({"params":model.coeff}) # assign new coefficient to the optimizer
#print(len(optimizer.param_groups))
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad() #initialize
if USECALI:
optimizer_cali = optim.LBFGS([model_cali.temperature],
lr=0.01,
max_iter=50)
optimizer_cali.zero_grad()
nll_criterion = BCEWithLogitsLoss().cuda() # BCE!!
ece_criterion = ECELoss().cuda()
# loss/ bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
interp = nn.Upsample(
size=(input_size[1], input_size[0]), mode='bilinear'
) # okay it automatically change to functional.interpolate
# 321, 321
if version.parse(torch.__version__) >= version.parse('0.4.0'): #0.4.1
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear')
# labels for adversarial training
pred_label = 0
gt_label = 1
semi_ratio_sum = 0
semi_sum = 0
loss_seg_sum = 0
loss_adv_sum = 0
loss_vat_sum = 0
l_seg_sum = 0
l_vat_sum = 0
l_adv_sum = 0
logits_list = []
labels_list = []
#https: // towardsdatascience.com / understanding - pytorch -with-an - example - a - step - by - step - tutorial - 81fc5f8c4e8e
for i_iter in range(args.num_steps):
loss_seg_value = 0 # L_seg
loss_adv_pred_value = 0 # 0.01 L_adv
loss_D_value = 0 # L_D
loss_semi_value = 0 # 0.1 L_semi
loss_semi_adv_value = 0 # 0.001 L_adv
loss_vat_value = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter) #changing lr by iteration
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(
): # <class 'torch.nn.parameter.Parameter'>, convolution weights
param.requires_grad = False # do not update gradient of D (freeze) while G
######### do unlabeled first!! 0.001 L_adv + 0.1 L_semi ###############
# lambda_semi, lambda_adv for unlabeled
if (args.lambda_semi > 0 or args.lambda_semi_adv > 0
) and i_iter >= args.semi_start_adv:
try:
_, batch = trainloader_remain_iter.next(
) #remain = unlabeled
print(trainloader_remain_iter.next())
except:
trainloader_remain_iter = enumerate(
trainloader_remain) # impose counters
_, batch = trainloader_remain_iter.next()
# only access to img
images, _, _, _ = batch # <class 'torch.Tensor'>
images = Variable(images).cuda(
args.gpu) # <class 'torch.Tensor'>
pred = interp(
model(images)) # S(X), pred <class 'torch.Tensor'>
pred_remain = pred.detach(
) #use detach() when attempting to remove a tensor from a computation graph, will be used for D
# https://discuss.pytorch.org/t/clone-and-detach-in-v0-4-0/16861
# The difference is that detach refers to only a given variable on which it's called.
# torch.no_grad affects all operations taking place within the with statement. >> for context,
# requires_grad is for tensor
# pred >> (8,21,321,321), L_adv
D_out = interp(
model_D(F.softmax(pred))
) # D(S(X)), confidence, 8,1,321,321, not detached, there was not dim
D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(
axis=1) # (8,321,321) 0~1
# 0.001 L_adv!!!!
ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(
np.bool) # no ignore_mask for unlabeled adv
loss_semi_adv = args.lambda_semi_adv * bce_loss(
D_out, make_D_label(gt_label,
ignore_mask_remain)) #gt_label =1,
# -log(D(S(X)))
loss_semi_adv = loss_semi_adv / args.iter_size #normalization
loss_semi_adv_value += loss_semi_adv.data.cpu().numpy(
) / args.lambda_semi_adv
##--- visualization, pred(8,21,321,321), D_out_sigmoid(8,321,321)
"""
if i_iter % 1000 == 0:
vpred = pred.transpose(1, 2).transpose(2, 3).contiguous() # (8,321,321,21)
vpred = vpred.view(-1, 21) # (8*321*321, 21)
vlogsx = F.log_softmax(vpred) # torch.Tensor
vsemi_gt = pred.data.cpu().numpy().argmax(axis=1)
vsemi_gt = Variable(torch.FloatTensor(vsemi_gt).long()).cuda(gpu)
vlogsx = vlogsx.gather(1, vsemi_gt.view(-1, 1))
sx = F.softmax(vpred).gather(1, vsemi_gt.view(-1, 1))
vD_out_sigmoid = Variable(torch.FloatTensor(D_out_sigmoid)).cuda(gpu).view(-1, 1)
vlogsx = (vlogsx*(2.5*vD_out_sigmoid+0.5))
vlogsx = -vlogsx.squeeze(dim=1)
sx = sx.squeeze(dim=1)
vD_out_sigmoid = vD_out_sigmoid.squeeze(dim=1)
dsx = vD_out_sigmoid.data.cpu().detach().numpy()
vlogsx = vlogsx.data.cpu().detach().numpy()
sx = sx.data.cpu().detach().numpy()
plt.clf()
plt.figure(figsize=(15, 5))
plt.subplot(131)
plt.ylim(0, 0.004)
plt.scatter(dsx, vlogsx, s = 0.1) # variable requires grad cannot call numpy >> detach
plt.xlabel('D(S(X))')
plt.ylabel('Loss_Semi per Pixel')
plt.subplot(132)
plt.scatter(dsx, vlogsx, s = 0.1) # variable requires grad cannot call numpy >> detach
plt.xlabel('D(S(X))')
plt.ylabel('Loss_Semi per Pixel')
plt.subplot(133)
plt.scatter(dsx, sx, s=0.1)
plt.xlabel('D(S(X))')
plt.ylabel('S(x)')
plt.savefig('/home/eungyo/AdvSemiSeg/plot/' + str(i_iter) + '.png')
"""
if args.lambda_semi <= 0 or i_iter < args.semi_start:
loss_semi_adv.backward()
loss_semi_value = 0
else:
semi_gt = pred.data.cpu().numpy().argmax(
axis=1
) # pred=S(X) ((8,21,321,321)), semi_gt is not one-hot, 8,321,321
#(8, 321, 321)
if not USECALI:
semi_ignore_mask = (
D_out_sigmoid < args.mask_T
) # both (8,321,321) 0~1threshold!, numpy
semi_gt[
semi_ignore_mask] = 255 # Yhat, ignore pixel becomes 255
semi_ratio = 1.0 - float(semi_ignore_mask.sum(
)) / semi_ignore_mask.size # ignored pixels / H*W
print('semi ratio: {:.4f}'.format(semi_ratio))
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
confidence = torch.FloatTensor(
D_out_sigmoid) ## added, only pred is on cuda
loss_semi = args.lambda_semi * weighted_loss_calc(
pred, semi_gt, args.gpu, confidence)
else:
semi_ratio = 1
semi_gt = (torch.FloatTensor(semi_gt)) # (8,321,321)
confidence = torch.FloatTensor(
F.sigmoid(
model_cali.temperature_scale(D_out.view(
-1))).data.cpu().numpy()) # (8*321*321,)
loss_semi = args.lambda_semi * calibrated_loss_calc(
pred, semi_gt, args.gpu, confidence, accuracies,
n_bin
) # L_semi = Yhat * log(S(X)) # loss_calc(pred, semi_gt, args.gpu)
# pred(8,21,321,321)
if semi_ratio != 0:
loss_semi = loss_semi / args.iter_size
loss_semi_value += loss_semi.data.cpu().numpy(
) / args.lambda_semi
if args.method == 'vatent' or args.method == 'vat':
#v_loss = vat_loss(model, images, pred, eps=args.epsilon[i]) # R_vadv
weighted_v_loss = weighted_vat_loss(
model,
images,
pred,
confidence,
eps=args.epsilon)
if args.method == 'vatent':
#v_loss += entropy_loss(pred) # R_cent (conditional entropy loss)
weighted_v_loss += weighted_entropy_loss(
pred, confidence)
v_loss = weighted_v_loss / args.iter_size
loss_vat_value += v_loss.data.cpu().numpy()
loss_semi_adv += args.alpha * v_loss
loss_vat_sum += loss_vat_value
if i_iter % 100 == 0 and sub_i == 4:
l_vat_sum = loss_vat_sum / 100
if i_iter == 0:
l_vat_sum = l_vat_sum * 100
loss_vat_sum = 0
loss_semi += loss_semi_adv
loss_semi.backward(
) # 0.001 L_adv + 0.1 L_semi, backward == back propagation
else:
loss_semi = None
loss_semi_adv = None
###########train with source (labeled data)############### L_ce + 0.01 * L_adv
try:
_, batch = trainloader_iter.next()
except:
trainloader_iter = enumerate(trainloader) # safe coding
_, batch = trainloader_iter.next() #counter, batch
images, labels, _, _ = batch # also get labels images(8,321,321)
images = Variable(images).cuda(args.gpu)
ignore_mask = (
labels.numpy() == 255
) # ignored pixels == 255 >> 1, yes ignored mask for labeled data
pred = interp(model(images)) # S(X), 8,21,321,321
loss_seg = loss_calc(pred, labels,
args.gpu) # -Y*logS(X)= L_ce, not detached
if USED:
softsx = F.softmax(pred, dim=1)
D_out = interp(model_D(softsx)) # D(S(X)), L_adv
loss_adv_pred = bce_loss(
D_out, make_D_label(
gt_label,
ignore_mask)) # both 8,1,321,321, gt_label = 1
# L_adv = -log(D(S(X)), make_D_label is all 1 except ignored_region
loss = loss_seg + args.lambda_adv_pred * loss_adv_pred
if USECALI:
if (args.lambda_semi > 0 or args.lambda_semi_adv > 0
) and i_iter >= args.semi_start_adv:
with torch.no_grad():
_, prediction = torch.max(softsx, 1)
labels_mask = (
(labels > 0) *
(labels != 255)) | (prediction.data.cpu() > 0)
labels = labels[labels_mask]
prediction = prediction[labels_mask]
fake_mask = (labels.data.cpu().numpy() !=
prediction.data.cpu().numpy())
real_label = make_conf_label(
1, fake_mask
) # (10*321*321, ) 0 or 1 (fake or real)
logits = D_out.squeeze(dim=1)
logits = logits[labels_mask]
logits_list.append(logits) # initialize
labels_list.append(real_label)
if (i_iter * args.iter_size * args.batch_size + sub_i +
1) % train_dataset_size == 0:
logits = torch.cat(logits_list).cuda(
) # overall 5000 images in val, #logits >> 5000,100, (1464*321*321,)
labels = torch.cat(labels_list).cuda()
before_temperature_nll = nll_criterion(
logits, labels).item() ####modify
before_temperature_ece, _, _ = ece_criterion(
logits, labels) # (1464*321*321,)
before_temperature_ece = before_temperature_ece.item(
)
print('Before temperature - NLL: %.3f, ECE: %.3f' %
(before_temperature_nll,
before_temperature_ece))
def eval():
loss_cali = nll_criterion(
model_cali.temperature_scale(logits),
labels)
loss_cali.backward()
return loss_cali
optimizer_cali.step(
eval) # just one backward >> not 50 iterations
after_temperature_nll = nll_criterion(
model_cali.temperature_scale(logits),
labels).item()
after_temperature_ece, accuracies, n_bin = ece_criterion(
model_cali.temperature_scale(logits), labels)
after_temperature_ece = after_temperature_ece.item(
)
print('Optimal temperature: %.3f' %
model_cali.temperature.item())
print(
'After temperature - NLL: %.3f, ECE: %.3f' %
(after_temperature_nll, after_temperature_ece))
logits_list = []
labels_list = []
else:
loss = loss_seg
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_sum += loss_seg / args.iter_size
if USED:
loss_adv_sum += loss_adv_pred
if i_iter % 100 == 0 and sub_i == 4:
l_seg_sum = loss_seg_sum / 100
if USED:
l_adv_sum = loss_adv_sum / 100
if i_iter == 0:
l_seg_sum = l_seg_sum * 100
l_adv_sum = l_adv_sum * 100
loss_seg_sum = 0
loss_adv_sum = 0
loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size
if USED:
loss_adv_pred_value += loss_adv_pred.data.cpu().numpy(
) / args.iter_size
##################### train D!!!###########################
###########################################################
# bring back requires_grad
if USED:
for param in model_D.parameters():
param.requires_grad = True # before False.
############# train with pred S(X)############# labeled + unlabeled
pred = pred.detach(
) #orginally only use labeled data, freeze S(X) when train D,
# We do train D with the unlabeled data. But the difference is quite small
if args.D_remain: #default true
pred = torch.cat(
(pred, pred_remain), 0
) # pred_remain(unlabeled S(x)) is detached 16,21,321,321
ignore_mask = np.concatenate(
(ignore_mask, ignore_mask_remain),
axis=0) # 16,321,321
D_out = interp(
model_D(F.softmax(pred, dim=1))
) # D(S(X)) 16,1,321,321 # softmax(pred,dim=1) for 0.4, not nessesary
loss_D = bce_loss(D_out,
make_D_label(pred_label,
ignore_mask)) # pred_label = 0
# -log(1-D(S(X)))
loss_D = loss_D / args.iter_size / 2 # iter_size = 1, /2 because there is G and D
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()
################## train with gt################### only labeled
#VOCGT and VOCdataset can be reduced to one dataset in this repo.
# get gt labels Y
#print "before train gt"
try:
print(trainloader_gt_iter.next()) # len 732
_, batch = trainloader_gt_iter.next()
except:
trainloader_gt_iter = enumerate(trainloader_gt)
_, batch = trainloader_gt_iter.next()
#print "train with gt?"
_, labels_gt, _, _ = batch
D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu) #one_hot
ignore_mask_gt = (labels_gt.numpy() == 255
) # same as ignore_mask (8,321,321)
#print "finish"
D_out = interp(model_D(D_gt_v)) # D(Y)
loss_D = bce_loss(D_out,
make_D_label(gt_label,
ignore_mask_gt)) # log(D(Y))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()
optimizer.step()
if USED:
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir)) #snapshot
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.6f}, loss_semi = {5:.6f}, loss_semi_adv = {6:.3f}, loss_vat = {7: .5f}'
.format(i_iter, args.num_steps, loss_seg_value,
loss_adv_pred_value, loss_D_value, loss_semi_value,
loss_semi_adv_value, loss_vat_value))
# L_ce L_adv for labeled L_D L_semi L_adv for unlabeled
#loss_adv should be inversely proportional to the loss_D if they are seeing the same data.
# loss_adv_p is essentially the inverse loss of loss_D. We expect them to achieve a good balance during the adversarial training
# loss_D is around 0.2-0.5 >> good
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + str(args.num_steps) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + str(args.num_steps) + '_D.pth'))
#torch.save(state, osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth.tar'))
#torch.save(state_D, osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth.tar'))
break
if i_iter % 100 == 0 and sub_i == 4: #loss_seg_value
wdata = "iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.6f}, loss_semi = {5:.8f}, loss_semi_adv = {6:.3f}, l_vat_sum = {7: .5f}, loss_label = {8: .4}\n".format(
i_iter, args.num_steps, l_seg_sum, l_adv_sum, loss_D_value,
loss_semi_value, loss_semi_adv_value, l_vat_sum,
l_seg_sum + 0.01 * l_adv_sum)
#wdata2 = "{0:8d} {1:s} {2:s} {3:s} {4:s} {5:s} {6:s} {7:s} {8:s}\n".format(i_iter,str(model.coeff[0])[8:14],str(model.coeff[1])[8:14],str(model.coeff[2])[8:14],str(model.coeff[3])[8:14],str(model.coeff[4])[8:14],str(model.coeff[5])[8:14],str(model.coeff[6])[8:14],str(model.coeff[7])[8:14])
if i_iter == 0:
f2 = open("/home/eungyo/AdvSemiSeg/snapshots/log.txt", 'w')
f2.write(wdata)
f2.close()
#f3 = open("/home/eungyo/AdvSemiSeg/snapshots/coeff.txt", 'w')
#f3.write(wdata2)
#f3.close()
else:
f1 = open("/home/eungyo/AdvSemiSeg/snapshots/log.txt", 'a')
f1.write(wdata)
f1.close()
#f4 = open("/home/eungyo/AdvSemiSeg/snapshots/coeff.txt", 'a')
#f4.write(wdata2)
#f4.close()
if i_iter % args.save_pred_every == 0 and i_iter != 0: # 5000
print('taking snapshot ...')
#state = {'epoch':i_iter, 'state_dict':model.state_dict(),'optim_dict':optimizer.state_dict()}
#state_D = {'epoch':i_iter, 'state_dict': model_D.state_dict(), 'optim_dict': optimizer_D.state_dict()}
#torch.save(state, osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth.tar'))
#torch.save(state_D, osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth.tar'))
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
# prepare
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
gpu = args.gpu
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
# create network
model = Res_Deeplab(num_classes=args.num_classes)
# load pretrained 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)
# only copy the params that exist in current model (caffe-like)
new_params = model.state_dict().copy()
for name, param in new_params.items():
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
print('copy {}'.format(name))
model.load_state_dict(new_params)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# init D
model_D = detector.FlawDetector(in_channels=24)
# model_D = FCDiscriminator(num_classes=args.num_classes)
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D.train()
model_D.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
train_dataset = VOCDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
if args.partial_data is None:
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
else:
#sample partial data
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print('loading train ids from {}'.format(args.partial_id))
else:
train_ids = [_ for _ in range(0, train_dataset_size)]
np.random.shuffle(train_ids)
pickle.dump(train_ids,
open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb'))
train_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
train_remain_sampler = data.sampler.SubsetRandomSampler(
train_ids[partial_size:])
train_gt_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_remain_sampler,
num_workers=3,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
sampler=train_gt_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
# loss/ bilinear upsampling
minimum_loss = detector.MinimumCriterion()
detector_loss = detector.FlawDetectorCriterion()
# bce_loss = BCEWithLogitsLoss2d()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear')
# labels for adversarial training
pred_label = 0
gt_label = 1
for i_iter in range(args.num_steps):
if i_iter > 0 and i_iter % 1000 == 0:
val(model, args.gpu)
model.train()
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_value = 0
loss_semi_adv_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
# do semi first
if (args.lambda_semi > 0 or args.lambda_semi_adv > 0
) and i_iter >= args.semi_start_adv:
try:
_, batch = trainloader_remain_iter.__next__()
except:
trainloader_remain_iter = enumerate(trainloader_remain)
_, batch = trainloader_remain_iter.__next__()
# only access to img
images, _, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred = interp(model(images))
pred_remain = pred.detach()
D_out = model_D(images, pred)
ignore_mask_remain = np.zeros(D_out.shape).astype(np.bool)
loss_semi_adv = args.lambda_semi_adv * minimum_loss(
D_out) # ke: SSL loss for unlabeled data
loss_semi_adv = loss_semi_adv
#loss_semi_adv.backward()
loss_semi_adv_value += loss_semi_adv.data.cpu().numpy(
) / args.iter_size
loss_semi_adv.backward()
loss_semi_value = 0
else:
loss_semi = None
loss_semi_adv = None
# train with source (labeled data)
try:
_, batch = trainloader_iter.__next__()
except:
trainloader_iter = enumerate(trainloader)
_, batch = trainloader_iter.__next__()
images, labels, _, _ = batch
images = Variable(images).cuda(args.gpu)
ignore_mask = (labels.numpy() == 255)
pred = interp(model(images))
loss_seg = loss_calc(pred, labels, args.gpu)
D_out = interp(model_D(images, pred))
loss_adv_pred = args.lambda_adv_pred * minimum_loss(
D_out) # ke: SSL loss for labeled data
loss = loss_seg + loss_adv_pred
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size
loss_adv_pred_value += loss_adv_pred.data.cpu().numpy(
) / args.iter_size
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred from labeled data
pred = pred.detach()
D_out = interp(model_D(images, pred))
detect_gt = detector.generate_flaw_detector_gt(
pred,
labels.view(labels.shape[0], 1, labels.shape[1],
labels.shape[2]).cuda(args.gpu), NUM_CLASSES,
IGNORE_LABEL)
loss_D = detector_loss(D_out, detect_gt)
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()
# # train with gt
# # get gt labels
# try:
# _, batch = trainloader_gt_iter.__next__()
# except:
# trainloader_gt_iter = enumerate(trainloader_gt)
# _, batch = trainloader_gt_iter.__next__()
# _, labels_gt, _, _ = batch
# D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu)
# ignore_mask_gt = (labels_gt.numpy() == 255)
# D_out = interp(model_D(D_gt_v))
# loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt))
# loss_D = loss_D/args.iter_size/2
# loss_D.backward()
# loss_D_value += loss_D.data.cpu().numpy()
optimizer.step()
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.6f}, loss_adv_l = {3:.6f}, loss_D = {4:.6f}, loss_semi = {5:.6f}, loss_adv_u = {6:.6f}'
.format(i_iter, args.num_steps, loss_seg_value,
loss_adv_pred_value, loss_D_value, loss_semi_value,
loss_semi_adv_value))
if i_iter >= args.num_steps - 1:
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + str(args.num_steps) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + str(args.num_steps) + '_D.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth'))
end = timeit.default_timer()
def main():
"""Create the model and start the training."""
w, h = map(int, args.input_size_source.split(','))
input_size_source = (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 = 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)
elif args.model == 'VGG':
model = DeeplabVGG(num_classes=args.num_classes,
pretrained=True,
vgg16_caffe_path=args.restore_from)
# saved_state_dict = torch.load(args.restore_from)
# model.load_state_dict(saved_state_dict)
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
#Discrimintator setting
model_D = FCDiscriminator(num_classes=args.num_classes)
model_D.train()
model_D.cuda(args.gpu)
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()
# labels for adversarial training
source_adv_label = 0
target_adv_label = 1
#Dataloader
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(GTA5DataSet(args.translated_data_dir,
args.data_list,
max_iters=args.num_steps *
args.iter_size * args.batch_size,
crop_size=input_size_source,
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)
style_trainloader = data.DataLoader(GTA5DataSet(
args.stylized_data_dir,
args.data_list,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_source,
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)
style_trainloader_iter = enumerate(style_trainloader)
if STAGE == 1:
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,
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)
else:
#Dataloader for self-training
targetloader = data.DataLoader(cityscapesDataSetLabel(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
mean=IMG_MEAN,
set=args.set,
label_folder='Path to generated pseudo labels'),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
interp = 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)
# load checkpoint
model, model_D, optimizer, start_iter = load_checkpoint(
model,
model_D,
optimizer,
filename=args.snapshot_dir + 'checkpoint_' + CHECKPOINT + '.pth.tar')
for i_iter in range(start_iter, args.num_steps):
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
#train segementation network
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# train with source
if STAGE == 1:
if i_iter % 2 == 0:
_, batch = next(trainloader_iter)
else:
_, batch = next(style_trainloader_iter)
else:
_, batch = next(trainloader_iter)
image_source, label, _, _ = batch
image_source = Variable(image_source).cuda(args.gpu)
pred_source = model(image_source)
pred_source = interp(pred_source)
loss_seg_source = loss_calc(pred_source, label, args.gpu)
loss_seg_source_value = loss_seg_source.item()
loss_seg_source.backward()
if STAGE == 2:
# train with target
_, batch = next(targetloader_iter)
image_target, target_label, _, _ = batch
image_target = Variable(image_target).cuda(args.gpu)
pred_target = model(image_target)
pred_target = interp_target(pred_target)
#target segmentation loss
loss_seg_target = loss_calc(pred_target,
target_label,
gpu=args.gpu)
loss_seg_target.backward()
# optimize
optimizer.step()
if STAGE == 1:
# train with target
_, batch = next(targetloader_iter)
image_target, _, _ = batch
image_target = Variable(image_target).cuda(args.gpu)
pred_target = model(image_target)
pred_target = interp_target(pred_target)
#output-level adversarial training
D_output_target = model_D(F.softmax(pred_target))
loss_adv = bce_loss(
D_output_target,
Variable(
torch.FloatTensor(D_output_target.data.size()).fill_(
source_adv_label)).cuda(args.gpu))
loss_adv = loss_adv * args.lambda_adv
loss_adv.backward()
#train discriminator
for param in model_D.parameters():
param.requires_grad = True
pred_source = pred_source.detach()
pred_target = pred_target.detach()
D_output_source = model_D(F.softmax(pred_source))
D_output_target = model_D(F.softmax(pred_target))
loss_D_source = bce_loss(
D_output_source,
Variable(
torch.FloatTensor(D_output_source.data.size()).fill_(
source_adv_label)).cuda(args.gpu))
loss_D_target = bce_loss(
D_output_target,
Variable(
torch.FloatTensor(D_output_target.data.size()).fill_(
target_adv_label)).cuda(args.gpu))
loss_D_source = loss_D_source / 2
loss_D_target = loss_D_target / 2
loss_D_source.backward()
loss_D_target.backward()
#optimize
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print('iter = {0:8d}/{1:8d}, loss_seg_source = {2:.5f}'.format(
i_iter, args.num_steps, loss_seg_source_value))
if i_iter % args.save_pred_every == 0:
print('taking snapshot ...')
state = {
'iter': i_iter,
'model': model.state_dict(),
'model_D': model_D.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(
state,
osp.join(args.snapshot_dir,
'checkpoint_' + str(i_iter) + '.pth.tar'))
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_D_' + str(i_iter) + '.pth'))
cityscapes_eval_dir = osp.join(args.cityscapes_eval_dir,
str(i_iter))
if not os.path.exists(cityscapes_eval_dir):
os.makedirs(cityscapes_eval_dir)
eval(osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth'),
cityscapes_eval_dir, i_iter)
iou19, iou13, iou = compute_mIoU(cityscapes_eval_dir, i_iter)
outputfile = open(args.output_file, 'a')
outputfile.write(
str(i_iter) + '\t' + str(iou19) + '\t' +
str(iou.replace('\n', ' ')) + '\n')
outputfile.close()
def train(log_file, arch, dataset, batch_size, iter_size, num_workers,
partial_data, partial_data_size, partial_id, ignore_label, crop_size,
eval_crop_size, is_training, learning_rate, learning_rate_d,
supervised, lambda_adv_pred, lambda_semi, lambda_semi_adv, mask_t,
semi_start, semi_start_adv, d_remain, momentum, not_restore_last,
num_steps, power, random_mirror, random_scale, random_seed,
restore_from, restore_from_d, eval_every, save_snapshot_every,
snapshot_dir, weight_decay, device):
settings = locals().copy()
import cv2
import torch
import torch.nn as nn
from torch.utils import data, model_zoo
import numpy as np
import pickle
import torch.optim as optim
import torch.nn.functional as F
import scipy.misc
import sys
import os
import os.path as osp
import pickle
from model.deeplab import Res_Deeplab
from model.unet import unet_resnet50
from model.deeplabv3 import resnet101_deeplabv3
from model.discriminator import FCDiscriminator
from utils.loss import CrossEntropy2d, BCEWithLogitsLoss2d
from utils.evaluation import EvaluatorIoU
from dataset.voc_dataset import VOCDataSet
import logger
torch_device = torch.device(device)
import time
if log_file != '' and log_file != 'none':
if os.path.exists(log_file):
print('Log file {} already exists; exiting...'.format(log_file))
return
with logger.LogFile(log_file if log_file != 'none' else None):
if dataset == 'pascal_aug':
ds = VOCDataSet(augmented_pascal=True)
elif dataset == 'pascal':
ds = VOCDataSet(augmented_pascal=False)
else:
print('Dataset {} not yet supported'.format(dataset))
return
print('Command: {}'.format(sys.argv[0]))
print('Arguments: {}'.format(' '.join(sys.argv[1:])))
print('Settings: {}'.format(', '.join([
'{}={}'.format(k, settings[k])
for k in sorted(list(settings.keys()))
])))
print('Loaded data')
def loss_calc(pred, label):
"""
This function returns cross entropy loss for semantic segmentation
"""
# out shape batch_size x channels x h x w -> batch_size x channels x h x w
# label shape h x w x 1 x batch_size -> batch_size x 1 x h x w
label = label.long().to(torch_device)
criterion = CrossEntropy2d()
return criterion(pred, label)
def lr_poly(base_lr, iter, max_iter, power):
return base_lr * ((1 - float(iter) / max_iter)**(power))
def adjust_learning_rate(optimizer, i_iter):
lr = lr_poly(learning_rate, i_iter, num_steps, power)
optimizer.param_groups[0]['lr'] = lr
if len(optimizer.param_groups) > 1:
optimizer.param_groups[1]['lr'] = lr * 10
def adjust_learning_rate_D(optimizer, i_iter):
lr = lr_poly(learning_rate_d, i_iter, num_steps, power)
optimizer.param_groups[0]['lr'] = lr
if len(optimizer.param_groups) > 1:
optimizer.param_groups[1]['lr'] = lr * 10
def one_hot(label):
label = label.numpy()
one_hot = np.zeros((label.shape[0], ds.num_classes, label.shape[1],
label.shape[2]),
dtype=label.dtype)
for i in range(ds.num_classes):
one_hot[:, i, ...] = (label == i)
#handle ignore labels
return torch.tensor(one_hot,
dtype=torch.float,
device=torch_device)
def make_D_label(label, ignore_mask):
ignore_mask = np.expand_dims(ignore_mask, axis=1)
D_label = np.ones(ignore_mask.shape) * label
D_label[ignore_mask] = ignore_label
D_label = torch.tensor(D_label,
dtype=torch.float,
device=torch_device)
return D_label
h, w = map(int, eval_crop_size.split(','))
eval_crop_size = (h, w)
h, w = map(int, crop_size.split(','))
crop_size = (h, w)
# create network
if arch == 'deeplab2':
model = Res_Deeplab(num_classes=ds.num_classes)
elif arch == 'unet_resnet50':
model = unet_resnet50(num_classes=ds.num_classes)
elif arch == 'resnet101_deeplabv3':
model = resnet101_deeplabv3(num_classes=ds.num_classes)
else:
print('Architecture {} not supported'.format(arch))
return
# load pretrained parameters
if restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(restore_from)
else:
saved_state_dict = torch.load(restore_from)
# only copy the params that exist in current model (caffe-like)
new_params = model.state_dict().copy()
for name, param in new_params.items():
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
model.load_state_dict(new_params)
model.train()
model = model.to(torch_device)
# init D
model_D = FCDiscriminator(num_classes=ds.num_classes)
if restore_from_d is not None:
model_D.load_state_dict(torch.load(restore_from_d))
model_D.train()
model_D = model_D.to(torch_device)
print('Built model')
if snapshot_dir is not None:
if not os.path.exists(snapshot_dir):
os.makedirs(snapshot_dir)
ds_train_xy = ds.train_xy(crop_size=crop_size,
scale=random_scale,
mirror=random_mirror,
range01=model.RANGE01,
mean=model.MEAN,
std=model.STD)
ds_train_y = ds.train_y(crop_size=crop_size,
scale=random_scale,
mirror=random_mirror,
range01=model.RANGE01,
mean=model.MEAN,
std=model.STD)
ds_val_xy = ds.val_xy(crop_size=eval_crop_size,
scale=False,
mirror=False,
range01=model.RANGE01,
mean=model.MEAN,
std=model.STD)
train_dataset_size = len(ds_train_xy)
if partial_data_size != -1:
if partial_data_size > partial_data_size:
print('partial-data-size > |train|: exiting')
return
if partial_data == 1.0 and (partial_data_size == -1 or
partial_data_size == train_dataset_size):
trainloader = data.DataLoader(ds_train_xy,
batch_size=batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
trainloader_gt = data.DataLoader(ds_train_y,
batch_size=batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
trainloader_remain = None
print('|train|={}'.format(train_dataset_size))
print('|val|={}'.format(len(ds_val_xy)))
else:
#sample partial data
if partial_data_size != -1:
partial_size = partial_data_size
else:
partial_size = int(partial_data * train_dataset_size)
if partial_id is not None:
train_ids = pickle.load(open(partial_id))
print('loading train ids from {}'.format(partial_id))
else:
rng = np.random.RandomState(random_seed)
train_ids = list(rng.permutation(train_dataset_size))
if snapshot_dir is not None:
pickle.dump(train_ids,
open(osp.join(snapshot_dir, 'train_id.pkl'), 'wb'))
print('|train supervised|={}'.format(partial_size))
print('|train unsupervised|={}'.format(train_dataset_size -
partial_size))
print('|val|={}'.format(len(ds_val_xy)))
print('supervised={}'.format(list(train_ids[:partial_size])))
train_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
train_remain_sampler = data.sampler.SubsetRandomSampler(
train_ids[partial_size:])
train_gt_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
trainloader = data.DataLoader(ds_train_xy,
batch_size=batch_size,
sampler=train_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain = data.DataLoader(ds_train_xy,
batch_size=batch_size,
sampler=train_remain_sampler,
num_workers=3,
pin_memory=True)
trainloader_gt = data.DataLoader(ds_train_y,
batch_size=batch_size,
sampler=train_gt_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
testloader = data.DataLoader(ds_val_xy,
batch_size=1,
shuffle=False,
pin_memory=True)
print('Data loaders ready')
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(learning_rate),
lr=learning_rate,
momentum=momentum,
weight_decay=weight_decay)
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=learning_rate_d,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
# loss/ bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
print('Built optimizer')
# labels for adversarial training
pred_label = 0
gt_label = 1
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_mask_accum = 0
loss_semi_value = 0
loss_semi_adv_value = 0
t1 = time.time()
print('Training for {} steps...'.format(num_steps))
for i_iter in range(num_steps + 1):
model.train()
model.freeze_batchnorm()
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(iter_size):
# train G
if not supervised:
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# do semi first
if not supervised and (lambda_semi > 0 or lambda_semi_adv > 0 ) and i_iter >= semi_start_adv and \
trainloader_remain is not None:
try:
_, batch = next(trainloader_remain_iter)
except:
trainloader_remain_iter = enumerate(trainloader_remain)
_, batch = next(trainloader_remain_iter)
# only access to img
images, _, _, _ = batch
images = images.float().to(torch_device)
pred = model(images)
pred_remain = pred.detach()
D_out = model_D(F.softmax(pred, dim=1))
D_out_sigmoid = F.sigmoid(
D_out).data.cpu().numpy().squeeze(axis=1)
ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(
np.bool)
loss_semi_adv = lambda_semi_adv * bce_loss(
D_out, make_D_label(gt_label, ignore_mask_remain))
loss_semi_adv = loss_semi_adv / iter_size
#loss_semi_adv.backward()
loss_semi_adv_value += float(
loss_semi_adv) / lambda_semi_adv
if lambda_semi <= 0 or i_iter < semi_start:
loss_semi_adv.backward()
loss_semi_value = 0
else:
# produce ignore mask
semi_ignore_mask = (D_out_sigmoid < mask_t)
semi_gt = pred.data.cpu().numpy().argmax(axis=1)
semi_gt[semi_ignore_mask] = ignore_label
semi_ratio = 1.0 - float(
semi_ignore_mask.sum()) / semi_ignore_mask.size
loss_semi_mask_accum += float(semi_ratio)
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
loss_semi = lambda_semi * loss_calc(pred, semi_gt)
loss_semi = loss_semi / iter_size
loss_semi_value += float(loss_semi) / lambda_semi
loss_semi += loss_semi_adv
loss_semi.backward()
else:
loss_semi = None
loss_semi_adv = None
# train with source
try:
_, batch = next(trainloader_iter)
except:
trainloader_iter = enumerate(trainloader)
_, batch = next(trainloader_iter)
images, labels, _, _ = batch
images = images.float().to(torch_device)
ignore_mask = (labels.numpy() == ignore_label)
pred = model(images)
loss_seg = loss_calc(pred, labels)
if supervised:
loss = loss_seg
else:
D_out = model_D(F.softmax(pred, dim=1))
loss_adv_pred = bce_loss(
D_out, make_D_label(gt_label, ignore_mask))
loss = loss_seg + lambda_adv_pred * loss_adv_pred
loss_adv_pred_value += float(loss_adv_pred) / iter_size
# proper normalization
loss = loss / iter_size
loss.backward()
loss_seg_value += float(loss_seg) / iter_size
if not supervised:
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
if d_remain:
pred = torch.cat((pred, pred_remain), 0)
ignore_mask = np.concatenate(
(ignore_mask, ignore_mask_remain), axis=0)
D_out = model_D(F.softmax(pred, dim=1))
loss_D = bce_loss(D_out,
make_D_label(pred_label, ignore_mask))
loss_D = loss_D / iter_size / 2
loss_D.backward()
loss_D_value += float(loss_D)
# train with gt
# get gt labels
try:
_, batch = next(trainloader_gt_iter)
except:
trainloader_gt_iter = enumerate(trainloader_gt)
_, batch = next(trainloader_gt_iter)
_, labels_gt, _, _ = batch
D_gt_v = one_hot(labels_gt)
ignore_mask_gt = (labels_gt.numpy() == ignore_label)
D_out = model_D(D_gt_v)
loss_D = bce_loss(D_out,
make_D_label(gt_label, ignore_mask_gt))
loss_D = loss_D / iter_size / 2
loss_D.backward()
loss_D_value += float(loss_D)
optimizer.step()
optimizer_D.step()
sys.stdout.write('.')
sys.stdout.flush()
if i_iter % eval_every == 0 and i_iter != 0:
model.eval()
with torch.no_grad():
evaluator = EvaluatorIoU(ds.num_classes)
for index, batch in enumerate(testloader):
image, label, size, name = batch
size = size[0].numpy()
image = image.float().to(torch_device)
output = model(image)
output = output.cpu().data[0].numpy()
output = output[:, :size[0], :size[1]]
gt = np.asarray(label[0].numpy()[:size[0], :size[1]],
dtype=np.int)
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2),
dtype=np.int)
evaluator.sample(gt, output, ignore_value=ignore_label)
sys.stdout.write('+')
sys.stdout.flush()
per_class_iou = evaluator.score()
mean_iou = per_class_iou.mean()
loss_seg_value /= eval_every
loss_adv_pred_value /= eval_every
loss_D_value /= eval_every
loss_semi_mask_accum /= eval_every
loss_semi_value /= eval_every
loss_semi_adv_value /= eval_every
sys.stdout.write('\n')
t2 = time.time()
print(
'iter = {:8d}/{:8d}, took {:.3f}s, loss_seg = {:.6f}, loss_adv_p = {:.6f}, loss_D = {:.6f}, loss_semi_mask_rate = {:.3%} loss_semi = {:.6f}, loss_semi_adv = {:.3f}'
.format(i_iter, num_steps, t2 - t1, loss_seg_value,
loss_adv_pred_value, loss_D_value,
loss_semi_mask_accum, loss_semi_value,
loss_semi_adv_value))
for i, (class_name,
iou) in enumerate(zip(ds.class_names, per_class_iou)):
print('class {:2d} {:12} IU {:.2f}'.format(
i, class_name, iou))
print('meanIOU: ' + str(mean_iou) + '\n')
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_value = 0
loss_semi_mask_accum = 0
loss_semi_adv_value = 0
t1 = t2
if snapshot_dir is not None and i_iter % save_snapshot_every == 0 and i_iter != 0:
print('taking snapshot ...')
torch.save(
model.state_dict(),
osp.join(snapshot_dir, 'VOC_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth'))
if snapshot_dir is not None:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(snapshot_dir, 'VOC_' + str(num_steps) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(snapshot_dir, 'VOC_' + str(num_steps) + '_D.pth'))
def main():
# LD build for summary
# training_summary = tf.summary.FileWriter(os.path.join(SUMMARY_DIR, 'train'))
# val_summary = tf.summary.FileWriter(os.path.join(SUMMARY_DIR, 'val'))
# dice_placeholder = tf.placeholder(tf.float32, [], name='dice')
# loss_placeholder = tf.placeholder(tf.float32, [], name='loss')
# # image_placeholder = tf.placeholder(tf.float32, [400*2, 400*2], name='image')
# # prediction_placeholder = tf.placeholder(tf.float32, [400*2, 400*2], name='prediction')
# tf.summary.scalar('dice', dice_placeholder)
# tf.summary.scalar('loss', loss_placeholder)
# # tf.summary.image('image', image_placeholder, max_outputs=1)
# # tf.summary.image('prediction', prediction_placeholder, max_outputs=1)
# summary_op = tf.summary.merge_all()
# config = tf.ConfigProto()
# config.gpu_options.allow_growth = True
# sess = tf.Session(config=config)
perfix_name = 'Liver'
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
gpu = args.gpu
# create network
model = Res_Deeplab(num_classes=args.num_classes)
# load pretrained 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)
# only copy the params that exist in current model (caffe-like)
new_params = model.state_dict().copy()
for name, param in new_params.items():
print(name)
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
print('copy {}'.format(name))
model.load_state_dict(new_params)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# LD delete
'''
# init D
model_D = FCDiscriminator(num_classes=args.num_classes)
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D.train()
model_D.cuda(args.gpu)
'''
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
# LD ADD start
from dataset.LiverDataset.liver_dataset import LiverDataset
user_name = 'ld'
validation_interval = 800
max_steps = 1000000000
batch_size = 5
n_neighboringslices = 1
input_size = 400
output_size = 400
slice_type = 'axial'
oversample = False
# reset_counter = args.reset_counter
label_of_interest = 1
label_required = 0
magic_number = 26.91
max_slice_tries_val = 0
max_slice_tries_train = 2
fuse_labels = True
apply_crop = False
train_data_dir = "/home/" + user_name + "/Documents/dataset/ISBI2017/media/nas/01_Datasets/CT/LITS/Training_Batch_2"
test_data_dir = "/home/" + user_name + "/Documents/dataset/ISBI2017/media/nas/01_Datasets/CT/LITS/Training_Batch_1"
train_dataset = LiverDataset(data_dir=train_data_dir,
slice_type=slice_type,
n_neighboringslices=n_neighboringslices,
input_size=input_size,
oversample=oversample,
label_of_interest=label_of_interest,
label_required=label_required,
max_slice_tries=max_slice_tries_train,
fuse_labels=fuse_labels,
apply_crop=apply_crop,
interval=validation_interval,
is_training=True,
batch_size=batch_size,
data_augmentation=True)
val_dataset = LiverDataset(data_dir=test_data_dir,
slice_type=slice_type,
n_neighboringslices=n_neighboringslices,
input_size=input_size,
oversample=oversample,
label_of_interest=label_of_interest,
label_required=label_required,
max_slice_tries=max_slice_tries_val,
fuse_labels=fuse_labels,
apply_crop=apply_crop,
interval=validation_interval,
is_training=False,
batch_size=batch_size)
# LD ADD end
# LD delete
'''
train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size,
scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size,
scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN)
if args.partial_data is None:
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True)
else:
#sample partial data
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print('loading train ids from {}'.format(args.partial_id))
else:
train_ids = range(train_dataset_size)
np.random.shuffle(train_ids)
pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb'))
train_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size])
train_remain_sampler = data.sampler.SubsetRandomSampler(train_ids[partial_size:])
train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size])
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True)
trainloader_remain = data.DataLoader(train_dataset,
batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
'''
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
# LD delete
'''
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9,0.99))
optimizer_D.zero_grad()
'''
# loss/ bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
interp = nn.Upsample(size=(input_size, input_size), mode='bilinear')
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(input_size, input_size),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(input_size, input_size), mode='bilinear')
# labels for adversarial training
pred_label = 0
gt_label = 1
loss_list = []
for i_iter in range(iter_start, args.num_steps):
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_value = 0
loss_semi_adv_value = 0
num_prediction = 0
num_ground_truth = 0
num_intersection = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
# LD delete
'''
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
'''
for sub_i in range(args.iter_size):
# train G
# LD delete
'''
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# do semi first
if (args.lambda_semi > 0 or args.lambda_semi_adv > 0 ) and i_iter >= args.semi_start_adv :
try:
_, batch = trainloader_remain_iter.next()
except:
trainloader_remain_iter = enumerate(trainloader_remain)
_, batch = trainloader_remain_iter.next()
# only access to img
images, _, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred = interp(model(images))
pred_remain = pred.detach()
D_out = interp(model_D(F.softmax(pred)))
D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(axis=1)
ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(np.bool)
loss_semi_adv = args.lambda_semi_adv * bce_loss(D_out, make_D_label(gt_label, ignore_mask_remain))
loss_semi_adv = loss_semi_adv/args.iter_size
#loss_semi_adv.backward()
loss_semi_adv_value += loss_semi_adv.data.cpu().numpy()[0]/args.lambda_semi_adv
if args.lambda_semi <= 0 or i_iter < args.semi_start:
loss_semi_adv.backward()
loss_semi_value = 0
else:
# produce ignore mask
semi_ignore_mask = (D_out_sigmoid < args.mask_T)
semi_gt = pred.data.cpu().numpy().argmax(axis=1)
semi_gt[semi_ignore_mask] = 255
semi_ratio = 1.0 - float(semi_ignore_mask.sum())/semi_ignore_mask.size
print('semi ratio: {:.4f}'.format(semi_ratio))
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
loss_semi = args.lambda_semi * loss_calc(pred, semi_gt, args.gpu)
loss_semi = loss_semi/args.iter_size
loss_semi_value += loss_semi.data.cpu().numpy()[0]/args.lambda_semi
loss_semi += loss_semi_adv
loss_semi.backward()
else:
loss_semi = None
loss_semi_adv = None
'''
# train with source
# LD delete
'''
try:
_, batch = trainloader_iter.next()
except:
trainloader_iter = enumerate(trainloader)
_, batch = trainloader_iter.next()
images, labels, _, _ = batch
images = Variable(images).cuda(args.gpu)
'''
batch_image, batch_label = train_dataset.get_next_batch()
batch_image = np.transpose(batch_image, axes=(0, 3, 1, 2))
batch_image = np.concatenate(
[batch_image, batch_image, batch_image], axis=1)
# print('Shape: ', np.shape(batch_image))
batch_image_torch = torch.Tensor(batch_image)
images = Variable(batch_image_torch).cuda(args.gpu)
# LD delete
# ignore_mask = (labels.numpy() == 255)
pred = interp(model(images))
pred_ny = pred.data.cpu().numpy()
pred_ny = np.transpose(pred_ny, axes=(0, 2, 3, 1))
pred_label_ny = np.squeeze(np.argmax(pred_ny, axis=3))
# prepare for dice
# print('Shape of gt is: ', np.shape(batch_label))
# print('Shape of pred is: ', np.shape(pred_ny))
# print('Shape of pred_label is: ', np.shape(pred_label_ny))
num_prediction += np.sum(np.asarray(pred_label_ny, np.uint8))
num_ground_truth += np.sum(np.asarray(batch_label >= 1, np.uint8))
num_intersection += np.sum(
np.asarray(
np.logical_and(batch_label >= 1, pred_label_ny >= 1),
np.uint8))
# num_intersection += np.sum(np.asarray(batch_label >= 1, np.uint8) == np.asarray(pred_label_ny, np.uint8))
loss_seg = loss_calc(pred, batch_label, args.gpu)
# LD delete
'''
D_out = interp(model_D(F.softmax(pred)))
loss_adv_pred = bce_loss(D_out, make_D_label(gt_label, ignore_mask))
loss = loss_seg + args.lambda_adv_pred * loss_adv_pred
'''
loss = loss_seg
# print('Loss is: ', loss)
# proper normalization
loss = loss / args.iter_size
loss.backward()
# print('Loss of numpy is: ', loss_seg.data.cpu().numpy())
# print('Loss of numpy of zero is: ', loss_seg.data.cpu().numpy())
loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size
loss_list.append(loss_seg_value)
# loss_adv_pred_value += loss_adv_pred.data.cpu().numpy()[0]/args.iter_size
# train D
# LD delete
'''
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
if args.D_remain:
pred = torch.cat((pred, pred_remain), 0)
ignore_mask = np.concatenate((ignore_mask,ignore_mask_remain), axis = 0)
D_out = interp(model_D(F.softmax(pred)))
loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask))
loss_D = loss_D/args.iter_size/2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()[0]
# train with gt
# get gt labels
try:
_, batch = trainloader_gt_iter.next()
except:
trainloader_gt_iter = enumerate(trainloader_gt)
_, batch = trainloader_gt_iter.next()
_, labels_gt, _, _ = batch
D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu)
ignore_mask_gt = (labels_gt.numpy() == 255)
D_out = interp(model_D(D_gt_v))
loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt))
loss_D = loss_D/args.iter_size/2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()[0]
'''
optimizer.step()
# optimizer_D.step()
dice = (2 * num_intersection + 1e-7) / (num_prediction +
num_ground_truth + 1e-7)
print('exp = {}'.format(args.snapshot_dir))
print('iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}'.format(
i_iter, args.num_steps, loss_seg_value))
print(
'dice: %.4f, num_prediction: %d, num_ground_truth: %d, num_intersection: %d'
% (dice, num_prediction, num_ground_truth, num_intersection))
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
perfix_name + str(args.num_steps) + '.pth'))
# torch.save(model_D.state_dict(), osp.join(args.snapshot_dir, perfix_name +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, perfix_name + str(i_iter)+'.pth'))
save_model(model, args.snapshot_dir, perfix_name, i_iter, 2)
# torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, perfix_name +str(i_iter)+'_D.pth'))
# if i_iter % UPDATE_TENSORBOARD_INTERVAL and i_iter != 0:
# # update tensorboard
# feed_dict = {
# dice_placeholder: dice,
# loss_placeholder: np.mean(loss_list)
# }
# summery_value = sess.run(summary_op, feed_dict)
# training_summary.add_summary(summery_value, i_iter)
# training_summary.flush()
#
# # for validation
# val_num_prediction = 0
# val_num_ground_truth = 0
# val_num_intersection = 0
# loss_list = []
#
# for _ in range(VAL_EXECUTE_TIMES):
# batch_image, batch_label = val_dataset.get_next_batch()
# batch_image = np.transpose(batch_image, axes=(0, 3, 1, 2))
# batch_image = np.concatenate([batch_image, batch_image, batch_image], axis=1)
# # print('Shape: ', np.shape(batch_image))
# batch_image_torch = torch.Tensor(batch_image)
# images = Variable(batch_image_torch).cuda(args.gpu)
#
# # LD delete
# # ignore_mask = (labels.numpy() == 255)
# pred = interp(model(images))
# pred_ny = pred.data.cpu().numpy()
# pred_ny = np.transpose(pred_ny, axes=(0, 2, 3, 1))
# pred_label_ny = np.squeeze(np.argmax(pred_ny, axis=3))
# val_num_prediction += np.sum(np.asarray(pred_label_ny, np.uint8))
# val_num_ground_truth += np.sum(np.asarray(batch_label >= 1, np.uint8))
# val_num_intersection += np.sum(np.asarray(np.logical_and(batch_label >= 1, pred_label_ny >= 1), np.uint8))
#
# loss_seg = loss_calc(pred, batch_label, args.gpu)
# loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size
# loss_list.append(loss_seg)
# dice = (2 * val_num_intersection + 1e-7) / (val_num_prediction + val_num_ground_truth + 1e-7)
# feed_dict = {
# dice_placeholder: dice,
# loss_placeholder: np.mean(loss_list)
# }
# summery_value = sess.run(summary_op, feed_dict)
# val_summary.add_summary(summery_value, i_iter)
# val_summary.flush()
# loss_list = []
training_summary.close()
val_summary.close()
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
h, w = list(map(int, args.input_size.split(','))) # 321, 321
input_size = (h, w)
cudnn.enabled = True
# create network
model = Res_Deeplab(num_classes=args.num_classes) # num_classes = 21
# load pretrained 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)
# only copy the params that exist in current model (caffe-like)
new_params = model.state_dict().copy()
for name, param in list(new_params.items()):
# print(name)
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
# print(('copy {}'.format(name)))
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) # num_classes = 21，全卷积判别模型
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D.train()
model_D.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
train_dataset = VOCDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
if args.partial_data is None: # 使用全部数据
trainloader = data.DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=True, # batch_size = 10
num_workers=5,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
else:
# sample partial data 部分数据
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print(('loading train ids from {}'.format(args.partial_id)))
else:
train_ids = list(range(train_dataset_size))
np.random.shuffle(train_ids)
pickle.dump(train_ids,
open(osp.join(args.snapshot_dir, 'train_id.pkl'),
'wb')) # 将train_ids写入train_id.pkl
train_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
train_remain_sampler = data.sampler.SubsetRandomSampler(
train_ids[partial_size:])
train_gt_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
trainloader = data.DataLoader(
train_dataset, # 数据集中采样输入
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_remain_sampler,
num_workers=3,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
sampler=train_gt_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
# loss/ bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
# interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear')
# labels for adversarial training
pred_label = 0
gt_label = 1
for i_iter in range(args.num_steps): # num_steps = 20000
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_value = 0
loss_semi_adv_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): # iter_size = 1
# train G
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# do semi first
if (args.lambda_semi > 0 or args.lambda_semi_adv > 0
) and i_iter >= args.semi_start_adv:
try:
_, batch = next(trainloader_remain_iter)
except:
trainloader_remain_iter = enumerate(trainloader_remain)
_, batch = next(trainloader_remain_iter)
# only access to img 无标签数据
images, _, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred = interp(model(images))
pred_remain = pred.detach() # 返回一个新的Variable，不具有grade
D_out = interp(model_D(F.softmax(pred)))
D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(
axis=1)
ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(
np.bool)
loss_semi_adv = args.lambda_semi_adv * bce_loss(
D_out, make_D_label(gt_label, ignore_mask_remain))
loss_semi_adv = loss_semi_adv / args.iter_size
# loss_semi_adv.backward()
# print('bug,', loss_semi_adv.data.cpu().numpy())
# loss_semi_adv_value += loss_semi_adv.data.cpu().numpy()[0]/args.lambda_semi_adv
loss_semi_adv_value += loss_semi_adv.data.cpu().numpy(
) / args.lambda_semi_adv
if args.lambda_semi <= 0 or i_iter < args.semi_start:
loss_semi_adv.backward()
loss_semi_value = 0
else:
# produce ignore mask
semi_ignore_mask = (D_out_sigmoid < args.mask_T
) # mask_T = 0.2，阈值
semi_gt = pred.data.cpu().numpy().argmax(
axis=1) # 返回维度为1上的最大值的下标
semi_gt[semi_ignore_mask] = 255
semi_ratio = 1.0 - float(semi_ignore_mask.sum(
)) / semi_ignore_mask.size # 被忽略的点占的比重
print(('semi ratio: {:.4f}'.format(semi_ratio)))
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
loss_semi = args.lambda_semi * loss_calc(
pred, semi_gt, args.gpu)
loss_semi = loss_semi / args.iter_size
# loss_semi_value += loss_semi.data.cpu().numpy()[0]/args.lambda_semi
loss_semi_value += loss_semi.data.cpu().numpy(
) / args.lambda_semi
loss_semi += loss_semi_adv
loss_semi.backward()
else:
loss_semi = None
loss_semi_adv = None
# train with source
try:
_, batch = next(trainloader_iter)
except:
trainloader_iter = enumerate(trainloader)
_, batch = next(trainloader_iter)
images, labels, _, _ = batch # 有标签数据
images = Variable(images).cuda(args.gpu)
ignore_mask = (labels.numpy() == 255)
pred = interp(model(images)) # interp上采样
loss_seg = loss_calc(pred, labels,
args.gpu) # 语义分割的cross entropy loss
D_out = interp(model_D(F.softmax(pred))) # 得到判别模型输出的判别图
loss_adv_pred = bce_loss(D_out,
make_D_label(gt_label, ignore_mask))
loss = loss_seg + args.lambda_adv_pred * loss_adv_pred
# proper normalization
loss = loss / args.iter_size
loss.backward()
# loss_seg_value += loss_seg.data.cpu().numpy()[0]/args.iter_size
# loss_adv_pred_value += loss_adv_pred.data.cpu().numpy()[0]/args.iter_size
loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size
loss_adv_pred_value += loss_adv_pred.data.cpu().numpy(
) / args.iter_size
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
if args.D_remain:
pred = torch.cat((pred, pred_remain), 0)
ignore_mask = np.concatenate((ignore_mask, ignore_mask_remain),
axis=0)
D_out = interp(model_D(F.softmax(pred)))
loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
# loss_D_value += loss_D.data.cpu().numpy()[0]
loss_D_value += loss_D.data.cpu().numpy()
# train with gt
# get gt labels
try:
_, batch = next(trainloader_gt_iter)
except:
trainloader_gt_iter = enumerate(trainloader_gt)
_, batch = next(trainloader_gt_iter)
_, labels_gt, _, _ = batch
D_gt_v = Variable(one_hot(labels_gt)).cuda(
args.gpu) # 每个类别一张label图，batch * class * h * w
ignore_mask_gt = (labels_gt.numpy() == 255)
D_out = interp(model_D(D_gt_v)) # ground_truth输入判别模型
loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
# loss_D_value += loss_D.data.cpu().numpy()[0]
loss_D_value += loss_D.data.cpu().numpy()
optimizer.step()
optimizer_D.step()
print(('exp = {}'.format(args.snapshot_dir)))
print((
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, '
'loss_semi = {5:.3f}, loss_semi_adv = {6:.3f}'.format(
i_iter, args.num_steps, loss_seg_value, loss_adv_pred_value,
loss_D_value, loss_semi_value, loss_semi_adv_value)))
if i_iter >= args.num_steps - 1: # num_step = 20000
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + str(args.num_steps) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + str(args.num_steps) + '_D.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0: # save_pred_every = 5000
print('taking snapshot ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
gpu = args.gpu
# Define network
model = Res_Deeplab(num_classes=args.num_classes)
# load pretrained 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)
# only copy the params that exist in current model (caffe-like)
new_params = model.state_dict().copy()
for name, param in new_params.items():
print(name)
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
print('copy {}'.format(name))
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)
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D.train()
model_D.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
train_dataset = VOCDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
if args.partial_data is None:
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
else:
#sample partial data
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print('loading train ids from {}'.format(args.partial_id))
else:
train_ids = list(range(train_dataset_size))
np.random.shuffle(train_ids)
pickle.dump(train_ids,
open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb'))
train_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
train_remain_sampler = data.sampler.SubsetRandomSampler(
train_ids[partial_size:])
train_gt_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_remain_sampler,
num_workers=3,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
sampler=train_gt_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
# loss/ bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
bce_loss_1 = BCEWithLogitsLoss2dd()
criterion = nn.L1Loss()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear')
# labels for adversarial training
pred_label = 0
gt_label = 1
for i_iter in range(args.num_steps):
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_value = 0
loss_semi_adv_value = 0
loss_laplacian = 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
# do semi first
if (args.lambda_semi > 0 or args.lambda_semi_adv > 0
) and i_iter >= args.semi_start_adv:
try:
_, batch = trainloader_remain_iter.__next__()
except:
trainloader_remain_iter = enumerate(trainloader_remain)
_, batch = trainloader_remain_iter.__next__()
# only access to img
images, _, _, _ = batch
images = images.cuda(args.gpu)
pred = interp(model(images))
pred_remain = pred.detach()
D_out = interp(model_D(F.softmax(pred)))
D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(
axis=1)
ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(
np.bool)
l, _, _ = bce_loss_1(
D_out, make_D_label(gt_label, ignore_mask_remain))
loss_semi_adv = args.lambda_semi_adv * -l
loss_semi_adv = loss_semi_adv / args.iter_size
#loss_semi_adv.backward()
loss_semi_adv_value += loss_semi_adv.data.cpu().numpy(
) / args.lambda_semi_adv
if args.lambda_semi <= 0 or i_iter < args.semi_start:
loss_semi_adv.backward()
loss_semi_value = 0
else:
# produce ignore mask
semi_ignore_mask = (D_out_sigmoid < args.mask_T)
semi_gt = pred.data.cpu().numpy().argmax(axis=1)
semi_gt[semi_ignore_mask] = 255
semi_ratio = 1.0 - float(
semi_ignore_mask.sum()) / semi_ignore_mask.size
print('semi ratio: {:.4f}'.format(semi_ratio))
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
loss_semi = args.lambda_semi * loss_calc(
pred, semi_gt, args.gpu)
loss_semi = loss_semi / args.iter_size
loss_semi_value += loss_semi.data.cpu().numpy(
) / args.lambda_semi
loss_semi += loss_semi_adv
loss_semi.backward()
else:
loss_semi = None
loss_semi_adv = None
# train with source
try:
_, batch = trainloader_iter.__next__()
except:
trainloader_iter = enumerate(trainloader)
_, batch = trainloader_iter.__next__()
images, labels, _, _ = batch
images = images.cuda(args.gpu)
ignore_mask = (labels.numpy() == 255)
pred = interp(model(images))
for i in range(10):
imagess = torch.zeros(321, 321).cuda()
for j in range(21):
try:
imagess += pred[i, j, :, :].reshape(321, 321)
except IndexError:
pass
try:
label = labels[i, :, :].reshape(321, 321).cuda()
except IndexError:
pass
imagess = torch.from_numpy(
cv2.Laplacian(imagess.cpu().detach().numpy(), -1)).cuda()
labell = torch.from_numpy(
cv2.Laplacian(label.cpu().detach().numpy(), -1)).cuda()
imagess = imagess.reshape(1, 1, 321, 321)
labell = labell.reshape(1, 1, 321, 321)
l, _, _ = bce_loss_1(imagess, labell)
loss_laplacian += l
loss_laplacian = -loss_laplacian
loss_seg = loss_calc(pred, labels, args.gpu)
D_out = interp(model_D(F.softmax(pred)))
l, _, _ = bce_loss_1(D_out, make_D_label(gt_label, ignore_mask))
loss_adv_pred = -l
loss = loss_seg + args.lambda_adv_pred * loss_adv_pred + loss_laplacian
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size
loss_adv_pred_value += loss_adv_pred.data.cpu().numpy(
) / args.iter_size
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
if args.D_remain:
pred = torch.cat((pred, pred_remain), 0)
ignore_mask = np.concatenate((ignore_mask, ignore_mask_remain),
axis=0)
D_out = interp(model_D(F.softmax(pred)))
l, p, t = bce_loss_1(D_out, make_D_label(pred_label, ignore_mask))
gradient_penalty = compute_gradient_penalty(
torch.ones_like(pred) * pred_label, pred, model_D)
loss_D = l + 0.05 * gradient_penalty
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()
# train with gt
# get gt labels
try:
_, batch = trainloader_gt_iter.__next__()
except:
trainloader_gt_iter = enumerate(trainloader_gt)
_, batch = trainloader_gt_iter.__next__()
_, labels_gt, _, _ = batch
D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu)
ignore_mask_gt = (labels_gt.numpy() == 255)
D_out = interp(model_D(D_gt_v))
l, p, t = bce_loss_1(D_out, make_D_label(gt_label, ignore_mask_gt))
gradient_penalty = compute_gradient_penalty(
torch.ones_like(D_gt_v) * gt_label, D_gt_v, model_D)
loss_D = -l + 0.05 * gradient_penalty
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()
optimizer.step()
if sub_i / 5 == 0:
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_semi = {5:.3f}, loss_semi_adv = {6:.3f}, loss_laplacian = {7:.3f}'
.format(i_iter, args.num_steps, loss_seg_value,
loss_adv_pred_value, loss_D_value, loss_semi_value,
loss_semi_adv_value, loss_laplacian))
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + str(args.num_steps) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + 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, 'VOC_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
# 将参数的input_size 映射到整数，并赋值，从字符串转换到整数二元组
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = False
gpu = args.gpu
# create network
model = Res_Deeplab(num_classes=args.num_classes)
# load pretrained 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)
# only copy the params that exist in current model (caffe-like)
# 确保模型中参数的格式与要加载的参数相同
# 返回一个字典，保存着module的所有状态（state）；parameters和persistent buffers都会包含在字典中，字典的key就是parameter和buffer的 names。
new_params = model.state_dict().copy()
for name, param in new_params.items():
# print (name)
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
# print('copy {}'.format(name))
model.load_state_dict(new_params)
# 设置为训练模式
model.train()
cudnn.benchmark = True
model.cuda(gpu)
# init D
model_D = FCDiscriminator(num_classes=args.num_classes)
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D.train()
model_D.cuda(gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
train_dataset = VOCDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
if args.partial_data is None:
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
else:
# sample partial data
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print('loading train ids from {}'.format(args.partial_id))
else:
train_ids = list(range(train_dataset_size)) # ?
np.random.shuffle(train_ids)
pickle.dump(train_ids,
open(osp.join(args.snapshot_dir, 'train_id.pkl'),
'wb')) # 写入文件
train_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
train_remain_sampler = data.sampler.SubsetRandomSampler(
train_ids[partial_size:])
train_gt_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_remain_sampler,
num_workers=3,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
sampler=train_gt_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
# loss/ bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear') # ???
# labels for adversarial training
pred_label = 0
gt_label = 1
for i_iter in range(args.num_steps):
print("Iter:", i_iter)
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_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
# do semi first
if args.lambda_semi > 0 and i_iter >= args.semi_start:
try:
_, batch = next(trainloader_remain_iter)
except:
trainloader_remain_iter = enumerate(trainloader_remain)
_, batch = next(trainloader_remain_iter)
# only access to img
images, _, _, _ = batch
images = Variable(images).cuda(gpu)
# images = Variable(images).cpu()
pred = interp(model(images))
D_out = interp(model_D(F.softmax(pred)))
D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(
axis=1)
# produce ignore mask
semi_ignore_mask = (D_out_sigmoid < args.mask_T)
semi_gt = pred.data.cpu().numpy().argmax(axis=1)
semi_gt[semi_ignore_mask] = 255
semi_ratio = 1.0 - float(
semi_ignore_mask.sum()) / semi_ignore_mask.size
print('semi ratio: {:.4f}'.format(semi_ratio))
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
loss_semi = args.lambda_semi * loss_calc(
pred, semi_gt, args.gpu)
loss_semi = loss_semi / args.iter_size
loss_semi.backward()
loss_semi_value += loss_semi.data.cpu().numpy(
)[0] / args.lambda_semi
else:
loss_semi = None
# train with source
try:
_, batch = next(trainloader_iter)
except:
trainloader_iter = enumerate(trainloader)
_, batch = next(trainloader_iter)
images, labels, _, _ = batch
images = Variable(images).cuda(gpu)
# images = Variable(images).cpu()
ignore_mask = (labels.numpy() == 255)
pred = interp(model(images))
loss_seg = loss_calc(pred, labels, args.gpu)
D_out = interp(model_D(F.softmax(pred)))
loss_adv_pred = bce_loss(D_out,
make_D_label(gt_label, ignore_mask))
loss = loss_seg + args.lambda_adv_pred * loss_adv_pred
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value += loss_seg.data.cpu().numpy()[0] / args.iter_size
loss_adv_pred_value += loss_adv_pred.data.cpu().numpy(
)[0] / args.iter_size
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
D_out = interp(model_D(F.softmax(pred)))
loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()[0]
# train with gt
# get gt labels
try:
_, batch = next(trainloader_gt_iter)
except:
trainloader_gt_iter = enumerate(trainloader_gt)
_, batch = next(trainloader_gt_iter)
_, labels_gt, _, _ = batch
D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu)
# D_gt_v = Variable(one_hot(labels_gt)).cpu()
ignore_mask_gt = (labels_gt.numpy() == 255)
D_out = interp(model_D(D_gt_v))
loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()[0]
optimizer.step()
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_semi = {5:.3f}'
.format(i_iter, args.num_steps, loss_seg_value,
loss_adv_pred_value, loss_D_value, loss_semi_value))
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + str(args.num_steps) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'VOC_' + 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, 'VOC_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
gpu = args.gpu
np.random.seed(args.random_seed)
# create network
model = Res_Deeplab(num_classes=args.num_classes)
# load pretrained 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)
# only copy the params that exist in current model (caffe-like)
new_params = model.state_dict().copy()
for name, param in new_params.items():
print(name)
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
print('copy {}'.format(name))
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)
# load dataset
train_dataset = VOCDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
if args.partial_data is None:
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
else:
#sample partial data
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print('loading train ids from {}'.format(args.partial_id))
else:
train_ids = np.arange(train_dataset_size)
np.random.shuffle(train_ids)
pickle.dump(train_ids,
open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb'))
# labeled data
train_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
train_gt_sampler = data.sampler.SubsetRandomSampler(
train_ids[:partial_size])
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_sampler,
num_workers=3,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
sampler=train_gt_sampler,
num_workers=3,
pin_memory=True)
# unlabeled data
train_remain_sampler = data.sampler.SubsetRandomSampler(
train_ids[partial_size:])
trainloader_remain = data.DataLoader(train_dataset,
batch_size=args.batch_size,
sampler=train_remain_sampler,
num_workers=3,
pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
# loss/bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear')
for i_iter in range(args.num_steps):
loss_seg_value = 0
loss_unlabeled_seg_value = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
for sub_i in range(args.iter_size):
# train Segmentation
# train with labeled images
try:
_, batch = trainloader_iter.next()
except:
trainloader_iter = enumerate(trainloader)
_, batch = trainloader_iter.next()
images, labels, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred = interp(model(images))
# computing loss
loss_seg = loss_calc(pred, labels, args.gpu)
# proper normalization
loss = loss_seg / args.iter_size
loss.backward()
loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size
# train with unlabeled
if args.lambda_semi > 0 and i_iter >= args.semi_start:
try:
_, batch = trainloader_remain_iter.next()
except:
trainloader_remain_iter = enumerate(trainloader_remain)
_, batch = trainloader_remain_iter.next()
# only access to img
images, _, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred = interp(model(images))
semi_gt = pred.data.cpu().numpy().argmax(axis=1)
semi_gt = torch.FloatTensor(semi_gt)
loss_unlabeled_seg = args.lambda_semi * loss_calc(
pred, semi_gt, args.gpu)
loss_unlabeled_seg = loss_unlabeled_seg / args.iter_size
loss_unlabeled_seg.backward()
loss_unlabeled_seg_value += loss_unlabeled_seg.data.cpu(
).numpy() / args.lambda_semi
else:
if args.lambda_semi > 0 and i_iter < args.semi_start:
loss_unlabeled_seg_value = 0
else:
loss_unlabeled_seg_value = None
optimizer.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_unlabeled_seg = {3:.3f} '
.format(i_iter, args.num_steps, loss_seg_value,
loss_unlabeled_seg_value))
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(
args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_' +
str(args.lambda_semi) + '_' + str(args.random_seed) +
'.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,
'VOC_' + str(i_iter) + '_' + str(args.lambda_semi) + '_' +
str(args.random_seed) + '.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')