def get_dataloader():
# TODO(xwd): Adaptive normalization by some large image.
# E.g. In medical image processing, WSI image is very large and different to ordinary images.
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
train_transform = transform.Compose([
transform.RandomScale([cfg['scale_min'], cfg['scale_max']]),
transform.RandomRotate([cfg['rotate_min'], cfg['rotate_max']],
padding=mean,
ignore_label=cfg['ignore_label']),
transform.RandomGaussianBlur(),
transform.RandomHorizontallyFlip(),
transform.RandomCrop([cfg['train_h'], cfg['train_w']],
crop_type='rand',
padding=mean,
ignore_label=cfg['ignore_label']),
transform.ToTensor(),
transform.Normalize(mean=mean, std=std)
])
train_data = cityscapes.Cityscapes(cfg['data_path'],
split='train',
transform=train_transform)
# Use data sampler to make sure each GPU loads specific parts of dataset to avoid data reduntant.
train_sampler = DistributedSampler(train_data)
train_loader = DataLoader(train_data,
batch_size=cfg['batch_size'] //
cfg['world_size'],
shuffle=(train_sampler is None),
num_workers=4,
pin_memory=True,
sampler=train_sampler,
drop_last=True)
return train_loader, train_sampler
def main(train_args):
print('No of classes', doc.num_classes)
if train_args['network'] == 'psp':
net = PSPNet(num_classes=doc.num_classes,
resnet=resnet,
res_path=res_path).cuda()
elif train_args['network'] == 'mfcn':
net = MFCN(num_classes=doc.num_classes, use_aux=True).cuda()
elif train_args['network'] == 'psppen':
net = PSPNet(num_classes=doc.num_classes,
resnet=resnet,
res_path=res_path).cuda()
print("number of cuda devices = ", torch.cuda.device_count())
if len(train_args['snapshot']) == 0:
curr_epoch = 1
train_args['best_record'] = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
else:
print('training resumes from ' + train_args['snapshot'])
net.load_state_dict(
torch.load(
os.path.join(ckpt_path, 'model_' + exp_name,
train_args['snapshot'])))
split_snapshot = train_args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
train_args['best_record'] = {
'epoch': int(split_snapshot[1]),
'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]),
'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]),
'fwavacc': float(split_snapshot[11])
}
net = torch.nn.DataParallel(net,
device_ids=list(
range(torch.cuda.device_count())))
net.train()
mean_std = ([0.9584, 0.9588, 0.9586], [0.1246, 0.1223, 0.1224])
weight = torch.FloatTensor(doc.num_classes)
#weight[0] = 1.0/0.5511 # background
train_simul_transform = simul_transforms.Compose([
simul_transforms.RandomSized(train_args['input_size']),
simul_transforms.RandomRotate(3),
#simul_transforms.RandomHorizontallyFlip()
simul_transforms.Scale(train_args['input_size']),
simul_transforms.CenterCrop(train_args['input_size'])
])
val_simul_transform = simul_transforms.Scale(train_args['input_size'])
train_input_transform = standard_transforms.Compose([
extended_transforms.RandomGaussianBlur(),
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
val_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
train_set = doc.DOC('train',
Dataroot,
joint_transform=train_simul_transform,
transform=train_input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=train_args['train_batch_size'],
num_workers=1,
shuffle=True,
drop_last=True)
train_loader_temp = DataLoader(train_set,
batch_size=1,
num_workers=1,
shuffle=True,
drop_last=True)
train_args['No_train_images'] = len(train_loader_temp)
del train_loader_temp
if train_args['No_train_images'] == 87677:
train_args['Type_of_train_image'] = 'All Synthetic slide image'
elif train_args['No_train_images'] == 150:
train_args['Type_of_train_image'] = 'Real image'
elif train_args['No_train_images'] == 84641:
train_args['Type_of_train_image'] = 'Synthetic image'
elif train_args['No_train_images'] == 151641:
train_args['Type_of_train_image'] = 'Real + Synthetic image'
val_set = doc.DOC('val',
Dataroot,
joint_transform=val_simul_transform,
transform=val_input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=1,
num_workers=1,
shuffle=False,
drop_last=True)
#criterion = CrossEntropyLoss2d(weight = weight, size_average = True, ignore_index = doc.ignore_label).cuda()
criterion = CrossEntropyLoss2d(size_average=True,
ignore_index=doc.ignore_label).cuda()
optimizer = optim.SGD([{
'params': [
param
for name, param in net.named_parameters() if name[-4:] == 'bias'
],
'lr':
2 * train_args['lr']
}, {
'params': [
param
for name, param in net.named_parameters() if name[-4:] != 'bias'
],
'lr':
train_args['lr'],
'weight_decay':
train_args['weight_decay']
}],
momentum=train_args['momentum'])
if len(train_args['snapshot']) > 0:
optimizer.load_state_dict(
torch.load(
os.path.join(ckpt_path, 'model_' + exp_name,
'opt_' + train_args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * train_args['lr']
optimizer.param_groups[1]['lr'] = train_args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(
os.path.join(ckpt_path, exp_name,
str(datetime.datetime.now()) + '.txt'),
'w').write(str(train_args) + '\n\n')
train(train_loader, net, criterion, optimizer, curr_epoch, train_args,
val_loader)
def main(train_args):
net = PSPNet(num_classes=voc.num_classes).cuda()
if len(train_args['snapshot']) == 0:
curr_epoch = 1
train_args['best_record'] = {'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0}
else:
print ('training resumes from ' + train_args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, train_args['snapshot'])))
split_snapshot = train_args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
train_args['best_record'] = {'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11])}
net.train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
train_simul_transform = simul_transforms.Compose([
simul_transforms.RandomSized(train_args['input_size']),
simul_transforms.RandomRotate(10),
simul_transforms.RandomHorizontallyFlip()
])
val_simul_transform = simul_transforms.Scale(train_args['input_size'])
train_input_transform = standard_transforms.Compose([
extended_transforms.RandomGaussianBlur(),
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
val_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.ToPILImage(),
])
visualize = standard_transforms.Compose([
standard_transforms.Scale(400),
standard_transforms.CenterCrop(400),
standard_transforms.ToTensor()
])
train_set = voc.VOC('train', joint_transform=train_simul_transform, transform=train_input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set, batch_size=train_args['train_batch_size'], num_workers=8, shuffle=True)
val_set = voc.VOC('val', joint_transform=val_simul_transform, transform=val_input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set, batch_size=1, num_workers=8, shuffle=False)
criterion = CrossEntropyLoss2d(size_average=True, ignore_index=voc.ignore_label).cuda()
optimizer = optim.SGD([
{'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
'lr': 2 * train_args['lr']},
{'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
'lr': train_args['lr'], 'weight_decay': train_args['weight_decay']}
], momentum=train_args['momentum'])
if len(train_args['snapshot']) > 0:
optimizer.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, 'opt_' + train_args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * train_args['lr']
optimizer.param_groups[1]['lr'] = train_args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(train_args) + '\n\n')
train(train_loader, net, criterion, optimizer, curr_epoch, train_args, val_loader, restore_transform, visualize)
def main():
"""Create the model and start the training."""
or_nyu_dict = {
0: 255,
1: 16,
2: 40,
3: 39,
4: 7,
5: 14,
6: 39,
7: 12,
8: 38,
9: 40,
10: 10,
11: 6,
12: 40,
13: 39,
14: 39,
15: 40,
16: 18,
17: 40,
18: 4,
19: 40,
20: 40,
21: 5,
22: 40,
23: 40,
24: 30,
25: 36,
26: 38,
27: 40,
28: 3,
29: 40,
30: 40,
31: 9,
32: 38,
33: 40,
34: 40,
35: 40,
36: 34,
37: 37,
38: 40,
39: 40,
40: 39,
41: 8,
42: 3,
43: 1,
44: 2,
45: 22
}
or_nyu_map = lambda x: or_nyu_dict.get(x, x) - 1
or_nyu_map = np.vectorize(or_nyu_map)
device = torch.device("cuda" if not args.cpu else "cpu")
w, h = map(int, args.input_size.split(','))
input_size = (w, h)
w, h = map(int, args.input_size_target.split(','))
input_size_target = (w, h)
cudnn.enabled = True
args.or_nyu_map = or_nyu_map
# 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)
elif args.restore_from == "":
saved_state_dict = None
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 == 40 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# print i_parts
model.load_state_dict(new_params)
model.train()
model.to(device)
cudnn.benchmark = True
if args.mode != "baseline" and args.mode != "baseline_tar":
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes).to(device)
model_D2 = FCDiscriminator(num_classes=args.num_classes).to(device)
model_D1.train()
model_D1.to(device)
model_D2.train()
model_D2.to(device)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
scale_min = 0.5
scale_max = 2.0
rotate_min = -10
rotate_max = 10
ignore_label = 255
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
args.width = w
args.height = h
train_transform = transforms.Compose([
# et.ExtResize( 512 ),
transforms.RandScale([scale_min, scale_max]),
transforms.RandRotate([rotate_min, rotate_max],
padding=IMG_MEAN_RGB,
ignore_label=ignore_label),
transforms.RandomGaussianBlur(),
transforms.RandomHorizontalFlip(),
transforms.Crop([args.height + 1, args.width + 1],
crop_type='rand',
padding=IMG_MEAN_RGB,
ignore_label=ignore_label),
#et.ExtRandomCrop(size=(opts.crop_size, opts.crop_size)),
#et.ExtColorJitter(brightness=0.5, contrast=0.5, saturation=0.5),
#et.ExtRandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=IMG_MEAN, std=[1, 1, 1]),
])
val_transform = transforms.Compose([
# et.ExtResize( 512 ),
transforms.Crop([args.height + 1, args.width + 1],
crop_type='center',
padding=IMG_MEAN_RGB,
ignore_label=ignore_label),
transforms.ToTensor(),
transforms.Normalize(mean=IMG_MEAN, std=[1, 1, 1]),
])
if args.mode != "baseline_tar":
src_train_dst = OpenRoomsSegmentation(root=args.data_dir,
opt=args,
split='train',
transform=train_transform,
imWidth=args.width,
imHeight=args.height,
remap_labels=args.or_nyu_map)
else:
src_train_dst = NYU_Labelled(root=args.data_dir_target,
opt=args,
split='train',
transform=train_transform,
imWidth=args.width,
imHeight=args.height,
phase="TRAIN",
randomize=True)
tar_train_dst = NYU(root=args.data_dir_target,
opt=args,
split='train',
transform=train_transform,
imWidth=args.width,
imHeight=args.height,
phase="TRAIN",
randomize=True,
mode=args.mode)
tar_val_dst = NYU(root=args.data_dir,
opt=args,
split='val',
transform=val_transform,
imWidth=args.width,
imHeight=args.height,
phase="TRAIN",
randomize=False)
trainloader = data.DataLoader(src_train_dst,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(tar_train_dst,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
if args.mode != "baseline" and args.mode != "baseline_tar":
optimizer_D1 = optim.Adam(model_D1.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D1.zero_grad()
optimizer_D2 = optim.Adam(model_D2.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D2.zero_grad()
if args.gan == 'Vanilla':
bce_loss = torch.nn.BCEWithLogitsLoss()
elif args.gan == 'LS':
bce_loss = torch.nn.MSELoss()
seg_loss = torch.nn.CrossEntropyLoss(ignore_index=255)
interp = nn.Upsample(size=(input_size[1] + 1, input_size[0] + 1),
mode='bilinear',
align_corners=True)
interp_target = nn.Upsample(size=(input_size_target[1] + 1,
input_size_target[0] + 1),
mode='bilinear',
align_corners=True)
# labels for adversarial training
source_label = 0
target_label = 1
# set up tensor board
if args.tensorboard:
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
for i_iter in range(args.num_steps):
loss_seg_value1 = 0
loss_seg_value1_tar = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_seg_value2 = 0
loss_seg_value2_tar = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
if args.mode != "baseline" and args.mode != "baseline_tar":
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate_D(optimizer_D1, i_iter)
adjust_learning_rate_D(optimizer_D2, i_iter)
sample_src = None
sample_tar = None
sample_res_src = None
sample_res_tar = None
sample_gt_src = None
sample_gt_tar = None
for sub_i in range(args.iter_size):
# train G
if args.mode != "baseline" and args.mode != "baseline_tar":
# 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:
trainloader_iter = enumerate(trainloader)
_, batch = trainloader_iter.__next__()
images, labels, _ = batch
sample_src = images.clone()
sample_gt_src = labels.clone()
images = images.to(device)
labels = labels.long().to(device)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
sample_pred_src = pred2.detach().cpu()
loss_seg1 = seg_loss(pred1, labels)
loss_seg2 = seg_loss(pred2, labels)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value1 += loss_seg1.item() / args.iter_size
loss_seg_value2 += loss_seg2.item() / args.iter_size
# train with target
try:
_, batch = targetloader_iter.__next__()
except:
targetloader_iter = enumerate(targetloader)
_, batch = targetloader_iter.__next__()
images, tar_labels, _, labelled = batch
n_labelled = labelled.sum().detach().item()
batch_size = images.shape[0]
sample_tar = images.clone()
sample_gt_tar = tar_labels.clone()
images = images.to(device)
pred_target1, pred_target2 = model(images)
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
#print("N_labelled {}".format(n_labelled))
if args.mode == "sda" and n_labelled != 0:
labelled = labelled.to(device) == 1
tar_labels = tar_labels.to(device)
loss_seg1_tar = seg_loss(pred_target1[labelled],
tar_labels[labelled])
loss_seg2_tar = seg_loss(pred_target2[labelled],
tar_labels[labelled])
loss_tar_labelled = loss_seg2_tar + args.lambda_seg * loss_seg1_tar
loss_tar_labelled = loss_tar_labelled / args.iter_size
loss_seg_value1_tar += loss_seg1_tar.item() / args.iter_size
loss_seg_value2_tar += loss_seg2_tar.item() / args.iter_size
else:
loss_tar_labelled = torch.zeros(
1, requires_grad=True).float().to(device)
# proper normalization
sample_pred_tar = pred_target2.detach().cpu()
if args.mode != "baseline" and args.mode != "baseline_tar":
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_adv_target1 = bce_loss(
D_out1,
torch.FloatTensor(
D_out1.data.size()).fill_(source_label).to(device))
loss_adv_target2 = bce_loss(
D_out2,
torch.FloatTensor(
D_out2.data.size()).fill_(source_label).to(device))
loss = args.lambda_adv_target1 * loss_adv_target1 + args.lambda_adv_target2 * loss_adv_target2
loss = loss / args.iter_size + loss_tar_labelled
#loss = loss_tar_labelled
loss.backward()
loss_adv_target_value1 += loss_adv_target1.item(
) / args.iter_size
loss_adv_target_value2 += loss_adv_target2.item(
) / args.iter_size
# train D
# bring back requires_grad
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
# train with source
pred1 = pred1.detach()
pred2 = pred2.detach()
D_out1 = model_D1(F.softmax(pred1))
D_out2 = model_D2(F.softmax(pred2))
loss_D1 = bce_loss(
D_out1,
torch.FloatTensor(
D_out1.data.size()).fill_(source_label).to(device))
loss_D2 = bce_loss(
D_out2,
torch.FloatTensor(
D_out2.data.size()).fill_(source_label).to(device))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
loss_D_value1 += loss_D1.item()
loss_D_value2 += loss_D2.item()
# train with target
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
D_out1 = model_D1(F.softmax(pred_target1))
D_out2 = model_D2(F.softmax(pred_target2))
loss_D1 = bce_loss(
D_out1,
torch.FloatTensor(
D_out1.data.size()).fill_(target_label).to(device))
loss_D2 = bce_loss(
D_out2,
torch.FloatTensor(
D_out2.data.size()).fill_(target_label).to(device))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
loss_D1.backward()
loss_D2.backward()
loss_D_value1 += loss_D1.item()
loss_D_value2 += loss_D2.item()
optimizer.step()
if args.mode != "baseline" and args.mode != "baseline_tar":
optimizer_D1.step()
optimizer_D2.step()
if args.tensorboard:
scalar_info = {
'loss_seg1': loss_seg_value1,
'loss_seg2': loss_seg_value2,
'loss_adv_target1': loss_adv_target_value1,
'loss_adv_target2': loss_adv_target_value2,
'loss_D1': loss_D_value1,
'loss_D2': loss_D_value2,
'loss_seg1_tar': loss_seg_value1_tar,
'loss_seg2_tar': loss_seg_value2_tar,
}
if i_iter % 10 == 0:
for key, val in scalar_info.items():
writer.add_scalar(key, val, i_iter)
if i_iter % 1000 == 0:
img = sample_src.cpu()[:, [2, 1, 0], :, :] + torch.from_numpy(
np.array(IMG_MEAN_RGB).reshape(1, 3, 1, 1)).float()
img = img.type(torch.uint8)
writer.add_images("Src/Images", img, i_iter)
label = tar_train_dst.decode_target(sample_gt_src).transpose(
0, 3, 1, 2)
writer.add_images("Src/Labels", label, i_iter)
preds = sample_pred_src.permute(0, 2, 3, 1).cpu().numpy()
preds = np.asarray(np.argmax(preds, axis=3), dtype=np.uint8)
preds = tar_train_dst.decode_target(preds).transpose(
0, 3, 1, 2)
writer.add_images("Src/Preds", preds, i_iter)
tar_img = sample_tar.cpu()[:,
[2, 1, 0], :, :] + torch.from_numpy(
np.array(IMG_MEAN_RGB).reshape(
1, 3, 1, 1)).float()
tar_img = tar_img.type(torch.uint8)
writer.add_images("Tar/Images", tar_img, i_iter)
tar_label = tar_train_dst.decode_target(
sample_gt_tar).transpose(0, 3, 1, 2)
writer.add_images("Tar/Labels", tar_label, i_iter)
tar_preds = sample_pred_tar.permute(0, 2, 3, 1).cpu().numpy()
tar_preds = np.asarray(np.argmax(tar_preds, axis=3),
dtype=np.uint8)
tar_preds = tar_train_dst.decode_target(tar_preds).transpose(
0, 3, 1, 2)
writer.add_images("Tar/Preds", tar_preds, i_iter)
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f} loss_adv1 = {4:.3f}, loss_adv2 = {5:.3f} loss_D1 = {6:.3f} loss_D2 = {7:.3f} loss_seg1_tar={8:.3f} loss_seg2_tar={9:.3f}'
.format(i_iter, args.num_steps, loss_seg_value1, loss_seg_value2,
loss_adv_target_value1, loss_adv_target_value2,
loss_D_value1, loss_D_value2, loss_seg_value1_tar,
loss_seg_value2_tar))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'OR_' + str(args.num_steps_stop) + '.pth'))
if args.mode != "baseline" and args.mode != "baseline_tar":
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'OR_' + str(args.num_steps_stop) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir,
'OR_' + 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, 'OR_' + str(i_iter) + '.pth'))
if args.mode != "baseline" and args.mode != "baseline_tar":
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'OR_' + str(i_iter) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir,
'OR_' + str(i_iter) + '_D2.pth'))
if args.tensorboard:
writer.close()