def make_data_loader(config):
train_set = spacenet.Spacenet(city=config.dataset,
split='train',
img_root=config.img_root,
source_dist=config.source_dist)
val_set = spacenet.Spacenet(city=config.dataset,
split='val',
img_root=config.img_root,
source_dist=config.source_dist)
num_class = train_set.NUM_CLASSES
train_loader = DataLoader(train_set,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.train_num_workers,
drop_last=True)
val_loader = DataLoader(val_set,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.val_num_workers,
drop_last=True)
test_loader = None
return train_loader, val_loader, test_loader, num_class
def __init__(self, model_path, source, target, cuda=False):
self.source_set = spacenet.Spacenet(city=source,
split='test',
img_root=config.img_root)
self.target_set = spacenet.Spacenet(city=target,
split='test',
img_root=config.img_root)
self.source_loader = DataLoader(self.source_set,
batch_size=16,
shuffle=False,
num_workers=2)
self.target_loader = DataLoader(self.target_set,
batch_size=16,
shuffle=False,
num_workers=2)
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path,
map_location=torch.device('cpu'))
#print(self.checkpoint.keys())
self.model.load_state_dict(self.checkpoint)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def make_data_loader(config):
train_set = spacenet.Spacenet(city=config.dataset,
split='train',
img_root=config.img_root,
gt_root=config.gt_root,
mean_std=config.mean_std,
if_augment=config.if_augment,
repeat_count=config.repeat_count,
transform_sample=config.transform_sample)
val_set = spacenet.Spacenet(city=config.dataset,
split='val',
img_root=config.img_root,
gt_root=config.gt_root,
mean_std=config.mean_std,
if_augment=config.if_augment,
repeat_count=config.repeat_count,
transform_sample=config.transform_sample)
#test_set = spacenet.Spacenet(city=config.dataset, split='test', img_root=config.img_root)
num_class = train_set.NUM_CLASSES
train_loader = DataLoader(train_set,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.train_num_workers,
drop_last=True)
val_loader = DataLoader(val_set,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.val_num_workers)
#test_loader = DataLoader(test_set, batch_size=config.batch_size, shuffle=False, num_workers=config.val_num_workers)
test_loader = None
#val_loader = None
return train_loader, val_loader, test_loader, num_class
def __init__(self, model_path, config, cuda=False):
self.target=config.all_dataset
self.target.remove(config.dataset)
# load source domain
self.source_set = spacenet.Spacenet(city=config.dataset, split='test', img_root=config.img_root)
self.source_loader = DataLoader(self.source_set, batch_size=16, shuffle=False, num_workers=2)
self.target_set = []
self.target_loader = []
# load other domains
for city in self.target:
tmp = spacenet.Spacenet(city=city, split='test', img_root=config.img_root)
self.target_set.append(tmp)
self.target_loader.append(DataLoader(tmp, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
#print(self.checkpoint.keys())
self.model.load_state_dict(self.checkpoint)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def compute_mean_variance(city='Vegas'):
#ini_dist = {'mean': (0.,0.,0.), 'std': (1., 1., 1.)}
train_set = spacenet.Spacenet(city=city,
split='train',
img_root='/usr/xtmp/satellite/spacenet/')
loader = DataLoader(train_set,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.train_num_workers,
drop_last=True)
fake_train_set = spacenet.Spacenet(
city=city,
split='train',
img_root='/usr/xtmp/satellite/FakeShanghai-Vegas/')
fake_loader = DataLoader(fake_train_set,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.train_num_workers,
drop_last=True)
mean = 0.
std = 0.
nb_samples = 0.
all_data = None
for sample in tqdm(loader):
data = sample['image']
batch_samples = data.size(0)
data = data.view(batch_samples, data.size(1), -1)
#if all_data is None:
# all_data = data
#else:
# torch.cat((all_data, data), 0)
mean += data.mean(2).sum(0)
std += data.std(2).sum(0)
nb_samples += batch_samples
for sample in tqdm(fake_loader):
data = sample['image']
batch_samples = data.size(0)
data = data.view(batch_samples, data.size(1), -1)
#if all_data is None:
# all_data = data
#else:
# torch.cat((all_data, data), 0)
mean += data.mean(2).sum(0)
std += data.std(2).sum(0)
nb_samples += batch_samples
mean /= nb_samples
std /= nb_samples
#mean = data.mean(2).sum(0)
#std = data.std(2).sum(0)
mean = mean.numpy().astype(float)
std = std.numpy().astype(float)
print('city: ', city)
print('mean: ', mean)
print('std:', std)
return (mean[0], mean[1], mean[2]), (std[0], std[1], std[2])
def __init__(self, model_path, config, bn, save_path, save_batch, sample_number, trial=100, cuda=False):
self.bn = bn
self.target=config.all_dataset
self.target.remove(config.dataset)
self.sample_number = sample_number
# load source domain
#self.source_set = spacenet.Spacenet(city=config.dataset, split='test', img_root=config.img_root, needs to be changed)
#self.source_loader = DataLoader(self.source_set, batch_size=16, shuffle=False, num_workers=2)
self.source_loader = None
self.save_path = save_path
self.save_batch = save_batch
self.trial = trial
self.target_set = []
self.target_loader = []
self.target_trainset = []
self.target_trainloader = []
self.config = config
# load other domains
for city in self.target:
test = spacenet.Spacenet(city=city, split='val', img_root=config.img_root, gt_root = config.gt_root, mean_std=config.mean_std, if_augment=config.if_augment, repeat_count=config.repeat_count)
self.target_set.append(test)
self.target_loader.append(DataLoader(test, batch_size=16, shuffle=False, num_workers=2))
train = spacenet.Spacenet(city=city, split='train', img_root=config.img_root, gt_root = config.gt_root, mean_std=config.mean_std, if_augment=config.if_augment, repeat_count=config.repeat_count, sample_number= sample_number)
self.target_trainset.append(train)
self.target_trainloader.append(DataLoader(train, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
#print(self.checkpoint.keys())
#self.model.load_state_dict(self.checkpoint)
self.model.load_state_dict(self.checkpoint['model'])
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def getImagesActivations():
images = []
activations = []
gts = []
cities = []
city_count = []
target = ['Vegas', 'Shanghai']
save_path = '/home/home1/swarnakr/main/DomainAdaptation/DA_viaBatchNorm/configs/ganDA/train_log/feats/'
for city in target:
train = spacenet.Spacenet(city=city,
split='train',
img_root='/usr/xtmp/satellite/spacenet/')
tbar = DataLoader(train, batch_size=16, shuffle=False, num_workers=2)
cc = 0
ct = 0
for i, sample in enumerate(tbar):
if True:
image, gt = sample['image'], sample['label']
activation = torch.load(
os.path.join(save_path, "{}Feats_{}.pt".format(
city[0], i))) #map_location=torch.device('cpu')))
assert image.shape[0] == activation.shape[0]
for kk in range(image.shape[0]):
images.append(image[kk].cpu().numpy())
activations.append(activation[kk].cpu().numpy())
gts.append(gt[kk].cpu().numpy())
cc += 1
city_count.append(len(images))
cities += [city] * city_count[ct]
ct += 1
return images, activations, gts, cities
def compute_mean_variance():
train_set = spacenet.Spacenet(city=config.dataset,
split='train',
img_root=config.img_root)
loader = DataLoader(train_set,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.train_num_workers,
drop_last=True)
mean = 0.
std = 0.
nb_samples = 0.
for sample in tqdm(loader):
data = sample['image']
batch_samples = data.size(0)
data = data.view(batch_samples, data.size(1), -1)
mean += data.mean(2).sum(0)
std += data.std(2).sum(0)
nb_samples += batch_samples
mean /= nb_samples
std /= nb_samples
print('mean: ', mean)
print('std:', std)
def getImagesActivations():
images = []
activations = []
gts = []
cities = []
city_count = []
target = ['Vegas', 'Shanghai']
save_path = '/home/home1/swarnakr/main/DomainAdaptation/DA_viaBatchNorm/configs/ganDA/train_log/feats/'
for city in target:
val = spacenet.Spacenet(city=city,
split='val',
img_root='/usr/xtmp/satellite/spacenet/')
tbar = DataLoader(val, batch_size=16, shuffle=False, num_workers=2)
cc = 0
ct = 0
for i, sample in enumerate(tbar):
if True:
image, gt = sample['image'], sample['label']
for kk in range(image.shape[0]):
images.append(image[kk].cpu().numpy())
gts.append(gt[kk].cpu().numpy())
cc += 1
city_count.append(len(images))
cities += [city] * city_count[ct]
ct += 1
return images, gts, cities
def __init__(self, model_path, config, bn, save_path, save_batch, cuda=False):
self.bn = bn
self.target=config.all_dataset
self.target.remove(config.dataset)
# load source domain
self.source_set = spacenet.Spacenet(city=config.dataset, split='test', img_root=config.img_root)
self.source_loader = DataLoader(self.source_set, batch_size=16, shuffle=False, num_workers=2)
self.save_path = save_path
self.save_batch = save_batch
self.target_set = []
self.target_loader = []
self.target_trainset = []
self.target_trainloader = []
self.config = config
# load other domains
for city in self.target:
#test_img_root = '/home/home1/swarnakr/main/DomainAdaptation/satellite/' + city + '/' + 'test'
#test = spacenet.Spacenet(city=city, split='test', img_root=test_img_root)
test = spacenet.Spacenet(city=city, split='test', img_root=config.img_root)
self.target_set.append(test)
self.target_loader.append(DataLoader(test, batch_size=16, shuffle=False, num_workers=2))
# train_img_root = '/home/home1/swarnakr/main/DomainAdaptation/satellite/' + city + '/' + 'train'
# train = spacenet.Spacenet(city=city, split='train', img_root=train_img_root)
train = spacenet.Spacenet(city=city, split='train', img_root=config.img_root)
self.target_trainset.append(train)
self.target_trainloader.append(DataLoader(train, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
#print(self.checkpoint.keys())
self.model.load_state_dict(self.checkpoint)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def make_data_loader(config):
train_set = spacenet.Spacenet(city=config.dataset,
split='train',
img_root=config.img_root)
val_set = spacenet.Spacenet(city=config.dataset,
split='val',
img_root=config.img_root)
#test_set = spacenet.Spacenet(city=config.dataset, split='test', img_root=config.img_root)
train_loader = DataLoader(train_set,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.train_num_workers,
drop_last=True)
val_loader = DataLoader(val_set,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.val_num_workers,
drop_last=True)
#test_loader = DataLoader(test_set, batch_size=config.batch_size, shuffle=False, num_workers=config.val_num_workers)
test_loader = None
return train_loader, val_loader, test_loader
def __init__(self,
model_path,
config,
bn,
save_path,
save_batch,
cuda=False):
self.bn = bn
self.city = config.dataset #all_dataset
self.save_path = save_path
self.save_batch = save_batch
self.target_trainset = []
self.target_trainloader = []
self.config = config
# load other domains
if 1: #for city in self.target:
train = spacenet.Spacenet(city=self.city,
split='train',
img_root=config.img_root)
self.target_trainset.append(train)
self.target_trainloader.append(
DataLoader(train, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
#if DA images
self.checkpoint = torch.load(model_path)
#'./train_log/' + self.config.dataset + '_da_' + city + '.pth')
self.model.load_state_dict(self.checkpoint)
if self.cuda:
self.model = self.model.cuda()
#grid_image = make_grid(torch.from_numpy(pred[:,None,:,:]), 5, normalize=False)
#self.writer.add_image(prefix+'/Prediction', grid_image, global_step)
#grid_image = make_grid(torch.from_numpy(target[:,None,:,:]), 5, normalize=False)
#self.writer.add_image(prefix+'/GT', grid_image, global_step)
grid_image = make_grid(torch.from_numpy(color_images(pred, target)),
5,
normalize=True,
range=(0, 255))
self.writer.add_image(prefix + '/Color', grid_image, global_step)
#grid_image = make_grid(images, 5, normalize=False)
#self.writer.add_image(prefix+'/Image', grid_image, global_step)
if __name__ == "__main__":
import sys
sys.path.append(os.getcwd())
from data import spacenet
from torch.utils.data import DataLoader
summary = TensorboardSummary('log')
dataset = spacenet.Spacenet('Vegas', source_dist=dist['Vegas'])
loader = DataLoader(dataset, batch_size=16, shuffle=False, num_workers=2)
for i, sample in enumerate(loader):
image, target = sample['image'], sample['label']
summary.visualize_image('train', 'Vegas', image, target,
image[:, :2, :, :], i)
if i == 3:
break
def main():
"""Create the model and start the training."""
w, h = map(int, args.input_size.split(','))
input_size = (w, h)
w, h = map(int, args.input_size_target.split(','))
input_size_target = (w, h)
cudnn.enabled = True
gpu = args.gpu
# Create network
if args.model == 'DeepLab':
model = DeeplabMulti(num_classes=args.num_classes)
# if args.restore_from[:4] == 'http' :
# saved_state_dict = model_zoo.load_url(args.restore_from)
# else:
# saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
# for i in saved_state_dict:
# # Scale.layer5.conv2d_list.3.weight
# i_parts = i.split('.')
# # print i_parts
# if not args.num_classes == 19 or not i_parts[1] == 'layer5':
# new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# # print i_parts
model.load_state_dict(new_params)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes)
model_D2 = FCDiscriminator(num_classes=args.num_classes)
model_D1.train()
model_D1.cuda(args.gpu)
model_D2.train()
model_D2.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
# trainloader = data.DataLoader(
# GTA5DataSet(args.data_dir, args.data_list, max_iters=args.num_steps * args.iter_size * args.batch_size,
# crop_size=input_size,
# scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN),
# batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True)
train_set = spacenet.Spacenet(city=config.dataset, split='train', img_root=config.img_root)
trainloader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True,
num_workers=args.num_workers, drop_last=True)
trainloader_iter = enumerate(trainloader)
# targetloader = data.DataLoader(cityscapesDataSet(args.data_dir_target, args.data_list_target,
# max_iters=args.num_steps * args.iter_size * args.batch_size,
# crop_size=input_size_target,
# scale=False, mirror=args.random_mirror, mean=IMG_MEAN,
# set=args.set),
# batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers,
# pin_memory=True)
target_set = spacenet.Spacenet(city=config.target, split='train', img_root=config.img_root)
targetloader = DataLoader(target_set, batch_size=args.batch_size, shuffle=True,
num_workers=args.num_workers, drop_last=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D1 = optim.Adam(model_D1.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
optimizer_D1.zero_grad()
optimizer_D2 = optim.Adam(model_D2.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
optimizer_D2.zero_grad()
if args.gan == 'Vanilla':
bce_loss = torch.nn.BCEWithLogitsLoss()
elif args.gan == 'LS':
bce_loss = torch.nn.MSELoss()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear')
# labels for adversarial training
source_label = 0
target_label = 1
for i_iter in range(args.num_steps):
loss_seg_value1 = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_seg_value2 = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate_D(optimizer_D1, i_iter)
adjust_learning_rate_D(optimizer_D2, i_iter)
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D1.parameters():
param.requires_grad = False
for param in model_D2.parameters():
param.requires_grad = False
# train with source
try:
_, batch = trainloader_iter.__next__()
except StopIteration:
trainloader_iter = enumerate(trainloader)
_, batch = trainloader_iter.__next__()
images, labels = batch
# print(images)
images = Variable(images).cuda(args.gpu)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
loss_seg1 = loss_calc(pred1, labels, args.gpu)
loss_seg2 = loss_calc(pred2, labels, args.gpu)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# proper normalization
loss = loss / args.iter_size
loss.backward()
# loss_seg_value1 += loss_seg1.data.cpu().numpy()[0] / args.iter_size
# loss_seg_value2 += loss_seg2.data.cpu().numpy()[0] / args.iter_size
loss_seg_value1 += loss_seg1.data.cpu().numpy() / args.iter_size
loss_seg_value2 += loss_seg2.data.cpu().numpy() / args.iter_size
# train with target
try:
_, batch = targetloader_iter.__next__()
except StopIteration:
targetloader_iter = enumerate(targetloader)
_, batch = targetloader_iter.__next__()
images, _ = batch
images = Variable(images).cuda(args.gpu)
pred_target1, pred_target2 = model(images)
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_adv_target1 = bce_loss(D_out1,
Variable(torch.FloatTensor(D_out1.data.size()).fill_(source_label)).cuda(
args.gpu))
loss_adv_target2 = bce_loss(D_out2,
Variable(torch.FloatTensor(D_out2.data.size()).fill_(source_label)).cuda(
args.gpu))
loss = args.lambda_adv_target1 * loss_adv_target1 + args.lambda_adv_target2 * loss_adv_target2
loss = loss / args.iter_size
loss.backward()
# loss_adv_target_value1 += loss_adv_target1.data.cpu().numpy()[0] / args.iter_size
# loss_adv_target_value2 += loss_adv_target2.data.cpu().numpy()[0] / args.iter_size
loss_adv_target_value1 += loss_adv_target1.data.cpu().numpy() / args.iter_size
loss_adv_target_value2 += loss_adv_target2.data.cpu().numpy() / args.iter_size
# train D
# bring back requires_grad
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
# train with source
pred1 = pred1.detach()
pred2 = pred2.detach()
D_out1 = model_D1(F.softmax(pred1))
D_out2 = model_D2(F.softmax(pred2))
loss_D1 = bce_loss(D_out1,
Variable(torch.FloatTensor(D_out1.data.size()).fill_(source_label)).cuda(args.gpu))
loss_D2 = bce_loss(D_out2,
Variable(torch.FloatTensor(D_out2.data.size()).fill_(source_label)).cuda(args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
loss_D_value1 += loss_D1.data.cpu().numpy()
loss_D_value2 += loss_D2.data.cpu().numpy()
# train with target
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_D1 = bce_loss(D_out1,
Variable(torch.FloatTensor(D_out1.data.size()).fill_(target_label)).cuda(args.gpu))
loss_D2 = bce_loss(D_out2,
Variable(torch.FloatTensor(D_out2.data.size()).fill_(target_label)).cuda(args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
loss_D_value1 += loss_D1.data.cpu().numpy()
loss_D_value2 += loss_D2.data.cpu().numpy()
optimizer.step()
optimizer_D1.step()
optimizer_D2.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f} loss_adv1 = {4:.3f}, loss_adv2 = {5:.3f} loss_D1 = {6:.3f} loss_D2 = {7:.3f}'.format(
i_iter, args.num_steps, loss_seg_value1, loss_seg_value2, loss_adv_target_value1, loss_adv_target_value2, loss_D_value1, loss_D_value2))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'paris_' + str(args.num_steps_stop) + '.pth'))
torch.save(model_D1.state_dict(), osp.join(args.snapshot_dir, 'paris_' + str(args.num_steps_stop) + '_D1.pth'))
torch.save(model_D2.state_dict(), osp.join(args.snapshot_dir, 'paris_' + str(args.num_steps_stop) + '_D2.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'paris_' + str(i_iter) + '.pth'))
torch.save(model_D1.state_dict(), osp.join(args.snapshot_dir, 'paris_' + str(i_iter) + '_D1.pth'))
torch.save(model_D2.state_dict(), osp.join(args.snapshot_dir, 'paris_' + str(i_iter) + '_D2.pth'))
