def __init__(self, args, model, train_loader, test_loader, log=None):
super().__init__()
self.args = args
self.trainingLoader = train_loader
self.testingLoader = test_loader
if (log is None):
self.log = Log(self.args)
else:
self.log = log
loss = TMGLowLoss(args, model).to(args.src_device)
self.parallel_loss = DataParallelCriterion(loss, args.device_ids)
def create_single_model(args):
print('Creating model with\n \
Input size: {}\n \
Output size: {}\n \
Activation {}\n \
Num layers: {}\n \
Hidden units per layer: {}\n \
Using bias: {}\n \
Using batchnorm {}\n \
With batchsize {}'.format( \
args.input_size, args.output_size, args.actv,
args.num_layers, args.hidden, args.bias, args.bn, args.batch))
model = args.model(input_size=args.input_size, output_size=args.output_size,
actv_type=args.actv,
num_layers=args.num_layers,
hidden_size=args.hidden,
bias=args.bias,
use_bn=args.bn)
args.loss = model.loss
if args.multi_gpu:
print('Using data parallelism with {} GPUs'.format(args.num_gpu))
#model = nn.DataParallel(model, device_ids = args.device_ids)
###
model = DataParallelModel(model, device_ids = args.device_ids)
args.loss = DataParallelCriterion(args.loss, device_ids = args.device_ids)
###
print('Sending model to device {}'.format(args.device))
model.to(args.device)
return model
def __init__(self, dataloader, hierarchical_transformer, config, i):
super(Trainer, self).__init__()
self.iter = i
self.config = config
self.cpu = torch.device("cpu")
self.multi_gpu = len(self.config.gpu_idx) > 1
self.dataloader = dataloader
self.word_encoder = WordEncoder.WordEncoder(config, self.dataloader.tweet_field.vocab)
self.word_pos_encoder = PositionEncoder.PositionEncoder(config, self.config.max_length)
self.time_delay_encoder = PositionEncoder.PositionEncoder(config, self.config.size)
# <----------- Check for GPU setting ----------->
if self.config.gpu:
self.hierarchical_transformer = DataParallelModel(hierarchical_transformer.cuda())
self.criterion = DataParallelCriterion(nn.NLLLoss())
else:
self.hierarchical_transformer = hierarchical_transformer
self.criterion = nn.NLLLoss()
self.adam_optimizer = optim.Adam(self.hierarchical_transformer.parameters(), np.power(self.config.d_model, - 0.5), betas = (self.config.beta_1, self.config.beta_2))
self.optimizer = Optimizer.Optimizer(self.config, self.adam_optimizer)
def init_fn(self, shared_model=None, **kwargs):
# Create auxiliary models
self.init_auxiliary()
if shared_model is not None:
self.model = shared_model
else:
self.model = self.init_model()
self.model = DataParallelModel(self.model.cuda(),
device_ids=self.gpus)
# self.model = torch.nn.DataParallel(self.model, device_ids=self.gpus).cuda()
# Setup a joint optimizer for the 2 models
self.optimizer = self.init_optimizer(self.options.optim.name)
self.lr_scheduler = self.init_lr(self.options.optim.lr_scheduler)
# Create loss functions
self.criterion = self.init_loss_functions()
self.criterion = DataParallelCriterion(self.criterion.cuda(),
device_ids=self.gpus)
# Create AverageMeters for losses
self.losses = AverageMeter()
# Evaluators
# self.evaluators = [Evaluator(self.options, self.logger, self.summary_writer, shared_model=self.model)]
self.dataset_size = None
def define_criterions(self):
"""Define criterions of losses
LSGAN loss for GAN losses, L1 loss for cycle, identity losses
Identity loss is used only for monet2photo task to keep the color context
If you had missed this in the paper, refer to section [5.2 - photo generation from paintings]
"""
self.criterionGAN = nn.MSELoss() # LSGAN losses
self.criterionCycle = nn.L1Loss()
self.criterionIdt = nn.L1Loss()
if opt.parallel and torch.cuda.device_count() > 1:
self.criterionGAN = DataParallelCriterion(self.criterionGAN)
self.criterionCycle = DataParallelCriterion(self.criterionCycle)
self.criterionIdt = DataParallelCriterion(self.criterionIdt)
def main(args):
# initialization
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
writer = SummaryWriter(log_dir=os.path.join(args.log_dir, args.method))
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.enabled = True
cudnn.benchmark = True
# conduct seg network
seg_model = get_model(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
new_params = seg_model.state_dict().copy()
if args.init:
for i in saved_state_dict:
i_parts = i.split('.')
if not i_parts[0] == 'fc':
new_params['encoder.' + '.'.join(i_parts[:])] = saved_state_dict[i]
seg_model.load_state_dict(new_params)
print('loading params w/o fc')
else:
seg_model.load_state_dict(saved_state_dict)
print('loading params all')
model = DataParallelModel(seg_model)
model.float()
model.cuda()
# define dataloader
train_loader = data.DataLoader(DataGenerator(root=args.root, list_path=args.lst,
crop_size=args.crop_size, training=True),
batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
val_loader = data.DataLoader(DataGenerator(root=args.val_root, list_path=args.val_lst,
crop_size=args.crop_size, training=False),
batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=True)
# define criterion & optimizer
criterion = ABRLovaszLoss(ignore_index=args.ignore_label, only_present=True, cls_p= args.num_classes, cls_h= args.hbody_cls, cls_f= args.fbody_cls)
criterion = DataParallelCriterion(criterion).cuda()
optimizer = optim.SGD(
[{'params': filter(lambda p: p.requires_grad, seg_model.parameters()), 'lr': args.learning_rate}],
lr=args.learning_rate, momentum=0.9, weight_decay=5e-4)
# key points
best_val_mIoU = 0
best_val_pixAcc = 0
start = time.time()
for epoch in range(0, args.epochs):
print('\n{} | {}'.format(epoch, args.epochs - 1))
# training
_ = train(model, train_loader, epoch, criterion, optimizer, writer)
# validation
if epoch %10 ==0 or epoch > args.epochs*0.8:
val_pixacc, val_miou = validation(model, val_loader, epoch, writer)
# save model
if val_pixacc > best_val_pixAcc:
best_val_pixAcc = val_pixacc
if val_miou > best_val_mIoU:
best_val_mIoU = val_miou
model_dir = os.path.join(args.snapshot_dir, args.method + '_miou.pth')
torch.save(seg_model.state_dict(), model_dir)
print('Model saved to %s' % model_dir)
os.rename(model_dir, os.path.join(args.snapshot_dir, args.method + '_miou'+str(best_val_mIoU)+'.pth'))
print('Complete using', time.time() - start, 'seconds')
print('Best pixAcc: {} | Best mIoU: {}'.format(best_val_pixAcc, best_val_mIoU))
def main():
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
random.seed(args.seed)
torch.manual_seed(args.seed)
writer = SummaryWriter(args.snapshot_dir)
os.environ["CUDA_VISIBLE_DEVICES"]=args.gpu
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
deeplab = get_segmentation_model("_".join([args.network, args.method]), num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
new_params = deeplab.state_dict().copy()
if 'wide' in args.network:
saved_state_dict = saved_state_dict['state_dict']
if 'vistas' in args.method:
saved_state_dict = saved_state_dict['body']
for i in saved_state_dict:
new_params[i] = saved_state_dict[i]
else:
for i in saved_state_dict:
i_parts = i.split('.')
if not 'classifier' in i_parts:
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
elif 'mobilenet' in args.network:
for i in saved_state_dict:
i_parts = i.split('.')
if not (i_parts[0]=='features' and i_parts[1]=='18') and not i_parts[0]=='classifier':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
else:
for i in saved_state_dict:
i_parts = i.split('.')
if not i_parts[0]=='fc' and not i_parts[0]=='last_linear' and not i_parts[0]=='classifier':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
if args.start_iters > 0:
deeplab.load_state_dict(saved_state_dict)
else:
deeplab.load_state_dict(new_params)
model = DataParallelModel(deeplab)
# model = nn.DataParallel(deeplab)
model.train()
model.float()
model.cuda()
criterion = CriterionCrossEntropy()
if "dsn" in args.method:
if args.ohem:
if args.ohem_single:
print('use ohem only for the second prediction map.')
criterion = CriterionOhemDSN_single(thres=args.ohem_thres, min_kept=args.ohem_keep, dsn_weight=float(args.dsn_weight))
else:
criterion = CriterionOhemDSN(thres=args.ohem_thres, min_kept=args.ohem_keep, dsn_weight=float(args.dsn_weight), use_weight=True)
else:
criterion = CriterionDSN(dsn_weight=float(args.dsn_weight), use_weight=True)
criterion = DataParallelCriterion(criterion)
criterion.cuda()
cudnn.benchmark = True
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(get_segmentation_dataset(args.dataset, root=args.data_dir, list_path=args.data_list,
max_iters=args.num_steps*args.batch_size, crop_size=input_size,
scale=args.random_scale, mirror=args.random_mirror, network=args.network),
batch_size=args.batch_size, shuffle=True, num_workers=1, pin_memory=True)
optimizer = optim.SGD([{'params': filter(lambda p: p.requires_grad, deeplab.parameters()), 'lr': args.learning_rate }],
lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)
optimizer.zero_grad()
for i_iter, batch in enumerate(trainloader):
sys.stdout.flush()
i_iter += args.start_iters
images, labels, _, _ = batch
images = Variable(images.cuda())
labels = Variable(labels.long().cuda())
optimizer.zero_grad()
lr = adjust_learning_rate(optimizer, i_iter)
if args.fix_lr:
lr = args.learning_rate
print('learning_rate: {}'.format(lr))
if 'gt' in args.method:
preds = model(images, labels)
else:
preds = model(images)
loss = criterion(preds, labels)
loss.backward()
optimizer.step()
if i_iter % 100 == 0:
writer.add_scalar('learning_rate', lr, i_iter)
writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter)
print('iter = {} of {} completed, loss = {}'.format(i_iter, args.num_steps, loss.data.cpu().numpy()))
if i_iter >= args.num_steps-1:
print('save model ...')
torch.save(deeplab.state_dict(),osp.join(args.snapshot_dir, 'CS_scenes_'+str(args.num_steps)+'.pth'))
break
if i_iter % args.save_pred_every == 0:
print('taking snapshot ...')
torch.save(deeplab.state_dict(),osp.join(args.snapshot_dir, 'CS_scenes_'+str(i_iter)+'.pth'))
end = timeit.default_timer()
print(end-start,'seconds')
def main(args):
# initialization
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
writer = SummaryWriter(log_dir=os.path.join(args.log_dir, args.method))
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.enabled = True
cudnn.benchmark = True
# conduct seg network
seg_model = get_model(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
new_params = seg_model.state_dict().copy()
# if args.init:
# for i in saved_state_dict:
# i_parts = i.split('.')
# if not i_parts[0] == 'fc':
# new_params['encoder.' + '.'.join(i_parts[:])] = saved_state_dict[i]
# seg_model.load_state_dict(new_params)
# print('loading params w/o fc')
# else:
# seg_model.load_state_dict(saved_state_dict)
# print('loading params all')
model = DataParallelModel(seg_model)
model.float()
model.cuda()
# define dataloader
train_loader = data.DataLoader(TrainGenerator(root=args.root,
list_path=args.lst,
crop_size=args.crop_size,
max_scale=2.0),
batch_size=args.batch_size,
shuffle=True,
num_workers=4,
pin_memory=True)
# define criterion & optimizer
criterion = ReportLovaszLoss(ignore_index=args.ignore_label,
only_present=True)
criterion = DataParallelCriterion(criterion).cuda()
optimizer = optim.SGD(
[{
'params': filter(lambda p: p.requires_grad,
seg_model.parameters()),
'lr': args.learning_rate
}],
lr=args.learning_rate,
momentum=0.9,
weight_decay=5e-4)
start = time.time()
for epoch in range(0, args.epochs):
print('\n{} | {}'.format(epoch, args.epochs - 1))
# training
_ = train(model, train_loader, epoch, criterion, optimizer, writer)
if epoch == args.epochs - 1:
model_dir = os.path.join(args.snapshot_dir,
args.method + '_final.pth')
torch.save(seg_model.state_dict(), model_dir)
print('Model saved to %s' % model_dir)
print('Complete using', time.time() - start, 'seconds')
saved_state_dict = torch.load(restore_from)
new_params = deeplab.state_dict().copy()
for i in saved_state_dict:
i_parts = i.split('.')
if not i_parts[0] == 'fc' and not i_parts[
0] == 'last_linear' and not i_parts[0] == 'classifier':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
deeplab.load_state_dict(new_params)
model = DataParallelModel(deeplab)
model.train()
model.float()
model.cuda()
criterion = CriterionCrossEntropy()
criterion = DataParallelCriterion(criterion)
criterion.cuda()
train_dataset = DatasetCityscapesAugmentation(root=data_dir,
list_path=data_list,
max_iters=num_steps *
batch_size,
crop_size=crop_size)
train_loader = data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1,
pin_memory=True)
optimizer = optim.SGD(
[{
class TrainFlow(object):
'''
Trains recursive encoder attention-driven decoder
Args:
args (argparse): object with programs arguments
model (torch.nn): Glow-TM model
train_loader (torch.dataloader): dataloader with training cases
test_loader (torch.dataloader): dataloader with training cases
log (Log): class for logging console outputs
'''
def __init__(self, args, model, train_loader, test_loader, log=None):
super().__init__()
self.args = args
self.trainingLoader = train_loader
self.testingLoader = test_loader
if (log is None):
self.log = Log(self.args)
else:
self.log = log
loss = TMGLowLoss(args, model).to(args.src_device)
self.parallel_loss = DataParallelCriterion(loss, args.device_ids)
def trainParallel(self, model, optimizer, tback=1, epoch=0, **kwargs):
'''
Trains the model for a single epoch
Args:
model (torch.nn.Module): PyTorch model to train
optimizer (torch.optim): PyTorch optimizer to update the models parameters
tback (int): number of time-steps to back propagate through in time
stride (int): The stride the low-fidelity input takes compared to output
epoch (int): current epoch
Returns:
total_loss (float): current loss
'''
model.train()
# Total training loss
total_loss = 0
beta = self.args.beta
print("Beta:", beta)
optimizer.zero_grad()
for mbIdx, (input0, target0,
lstm_seeds) in enumerate(self.trainingLoader):
aKey = model.module.initLSTMStates(
lstm_seeds,
[target0.size(-2), target0.size(-1)])
# aKey = model.module.initLSTMStates(torch.LongTensor(input0.size(0)).random_(0, int(1e8)), [target0.size(-2), target0.size(-1)])
a0 = copy.deepcopy(aKey)
loss = 0 # Time-series loss
# Loop of time-steps
tmax = target0.size(1)
tback = 10
input_next = input0[:, :tback].to(self.args.device)
target_next = target0[:, :tback].to(self.args.device)
target0_mean = torch.mean(target0, axis=1).to(self.args.device)
target0_rms = torch.sqrt(
torch.mean((target0.to(self.args.device) -
target0_mean.unsqueeze(1))**2,
dim=1)).to(self.args.device)
# Splits time-series into smaller blocks to calculate back-prop through time
for i in range(0, tmax // tback):
input = input_next
ytarget = target_next
# Asynch load the next time-series
if (i + 1 < tmax // tback):
input_next = input0[:, (i + 1) * tback:(i + 2) *
tback].cuda(self.args.device,
non_blocking=True)
target_next = target0[:, (i + 1) * tback:(i + 2) *
tback].cuda(self.args.device,
non_blocking=True)
loss = 0
gpu_loss = [0 for i in range(self.args.n_gpu)]
modelPredictions = TMGLowPredictionItem()
model.scatterModel()
for tstep in range(tback):
# Model forward
outputs = model.sample(input[:, tstep], a0)
if (isinstance(outputs, list)):
yPred = [output[0] for output in outputs]
logp = [output[1] for output in outputs]
a0 = [output[2] for output in outputs]
else:
yPred, logp, a0 = outputs
modelPredictions.add(yPred, logp, ytarget[:, tstep])
# Recompile recurrent states onto the source device
if (self.args.n_gpu > 1):
a0 = model.gatherLSTMStates(a0)
else:
a0 = outputs[2]
# Compute the reverse KL divergence loss
outputs = modelPredictions.getOutputs()
targets = modelPredictions.getTargets()
loss0 = self.parallel_loss(outputs, targets, target0_mean,
target0_rms)
if (self.args.n_gpu > 1):
gpu_loss = [
gpu_loss[i] + loss0[i] for i in range(len(loss0))
]
else:
gpu_loss = [gpu_loss[0] + loss0]
modelPredictions.clear()
loss = self.parallel_loss.gather(
gpu_loss, output_device=self.args.src_device).mean()
# Backwards!
loss.backward()
# print(getGpuMemoryMap())
torch.nn.utils.clip_grad_norm_(model.parameters(),
self.args.max_grad_norm)
optimizer.step()
optimizer.zero_grad()
# Average the LSTM states with the initial state to prevent convergence
for j in range(len(a0)):
a_out, c_out = a0[j]
a_key, c_key = aKey[j]
a0[j] = (0.5 * a_out.detach() + 0.5 * a_key,
0.5 * c_out.detach() + 0.5 * c_key)
total_loss = total_loss + loss.detach()
# Sync cuda processes here
# Note sure if needed, but hopefully makes sure next data is loaded.
torch.cuda.synchronize()
torch.cuda.empty_cache()
# Add loss of time-series to total loss
# Mini-batch progress log
if ((mbIdx + 1) % 5 == 0):
self.log.log(
'Train Epoch: {}; Mini-batch: {}/{} ({:.0f}%); \t Current Loss: {:.6f}'
.format(epoch, mbIdx, len(self.trainingLoader),
100. * mbIdx / len(self.trainingLoader),
total_loss))
return total_loss
def test(self, model, samples=1, epoch=0, plot=True):
'''
Tests the model
Args:
model (torch.nn.Module): PyTorch model to test
stride (int): The stride the low-fidelity input takes compared to output
samples (int): Number of prediction to sample from the model
epoch (int): current epoch
plot (boolean): If to plot two of the predictions or not
Returns:
mse (float): mean-squared-error between the predictive mean and targets
'''
model.eval()
# Total test loss
total_loss = 0
out_std = model.module.out_std.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
out_mu = model.module.out_mu.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
for mbIdx, (input0, target0, u0) in enumerate(self.testingLoader):
u0 = u0.to(self.args.device).unsqueeze(-1).unsqueeze(-1).unsqueeze(
-1).unsqueeze(-1)
u0 = torch.cat((u0 / 2.0, u0 / 2.0, u0**2), dim=2)
u0in = torch.ones(input0[:, :, :1].size()).to(
self.args.device) * u0[:, :, 0, :, :].unsqueeze(2)
# input = torch.cat((input0.to(self.args.device), u0in), dim=2)
input = input0.to(self.args.device)
ytarget = out_std * target0.to(self.args.device) + out_mu
dims = [samples] + list(ytarget.size())
yPred = torch.zeros(dims).type(ytarget.type())
# Max number of time steps
tmax = 40
# Loop through samples
for i in range(samples):
aKey = model.module.initLSTMStates(
torch.LongTensor(input.size(0)).random_(0, int(1e8)),
[ytarget.size(-2), ytarget.size(-1)])
a0 = copy.deepcopy(aKey)
# Loop of time-steps
model.scatterModel()
for tstep in range(0, tmax + 1):
# Model forward
outputs = model.sample(input[:, tstep], a0)
yPred0, logp, a0 = model.gather(outputs,
self.args.src_device)
out_std = model.module.out_std.unsqueeze(0).unsqueeze(
-1).unsqueeze(-1)
out_mu = model.module.out_mu.unsqueeze(0).unsqueeze(
-1).unsqueeze(-1)
yPredHat = out_std * yPred0 + out_mu
yPred[i, :, tstep] = yPredHat.detach()
# Average current LSTM states with initial state
if (tstep % 10 == 0):
for j in range(len(a0)):
a_out, c_out = a0[j]
a_key, c_key = aKey[j]
a0[j] = (0.5 * a_out.detach() + 0.5 * a_key,
0.5 * c_out.detach() + 0.5 * c_key)
if (plot and mbIdx == 0):
self.log.log('Plotting predictions.')
plotVelocityPred(self.args,
input,
yPred,
ytarget,
bidx=0,
stride=4,
epoch=epoch)
plotVelocityPred(self.args,
input,
yPred,
ytarget,
bidx=1,
stride=4,
epoch=epoch)
# Summation of the squared error between the mean of the samples and target
total_loss = total_loss + (torch.pow(
torch.mean(yPred[:, :, 1:tmax + 1], dim=0) -
ytarget[:, 1:tmax + 1], 2)).sum().detach()
# Return the mse
return total_loss / (self.args.ntest * tmax * yPred.size(-2) *
yPred.size(-1))
class CycleGAN:
"""GANs which use cycle loss to make generators to map the image
using meaningful connections between domains, which keeps content of the image
so that generators can train how to do style-transfer without paired dataset.
"""
def __init__(self,
lr=2e-4,
betas=(.5, .999),
n_epochs=200,
ngf=64,
ndf=64,
lambdaA=10.,
lambdaB=10.,
lambdaIdt=.5):
self.learning_rate = lr
self.betas = betas
self.n_epochs = n_epochs
self.ngf = ngf
self.ndf = ndf
self.lambdaA = lambdaA # coefficients of cycle losses
self.lambdaB = lambdaB
self.lambdaIdt = lambdaIdt # coefficient of Identity losses
self.true_label = torch.tensor(1.)
self.false_label = torch.tensor(0.)
# todo: need to add data parallel code
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.fakeA_pool = ImagePool(opt.image_pool_size)
self.fakeB_pool = ImagePool(opt.image_pool_size)
self.define_nets()
self.define_criterions()
self.move_to()
self.define_optimizers()
self.define_schedulers()
def define_nets(self):
"""Define generators and discriminators for both directions"""
self.netG_A = define_G(self.ngf)
self.netG_B = define_G(self.ngf)
self.netD_A = define_D(self.ndf)
self.netD_B = define_D(self.ndf)
self.G_params = list(self.netG_A.parameters()) + list(
self.netG_B.parameters())
self.D_params = list(self.netD_A.parameters()) + list(
self.netD_B.parameters())
def define_criterions(self):
"""Define criterions of losses
LSGAN loss for GAN losses, L1 loss for cycle, identity losses
Identity loss is used only for monet2photo task to keep the color context
If you had missed this in the paper, refer to section [5.2 - photo generation from paintings]
"""
self.criterionGAN = nn.MSELoss() # LSGAN losses
self.criterionCycle = nn.L1Loss()
self.criterionIdt = nn.L1Loss()
if opt.parallel and torch.cuda.device_count() > 1:
self.criterionGAN = DataParallelCriterion(self.criterionGAN)
self.criterionCycle = DataParallelCriterion(self.criterionCycle)
self.criterionIdt = DataParallelCriterion(self.criterionIdt)
def define_optimizers(self):
"""Define optimizers"""
self.optimizerG = optim.Adam(self.G_params,
lr=self.learning_rate,
betas=self.betas)
self.optimizerD = optim.Adam(self.D_params,
lr=self.learning_rate,
betas=self.betas)
def define_schedulers(self):
"""Define schedulers
for <100 epoch, maintain initial learning rate
and for >=100 epoch, linearly decay to 0"""
def lambda_rule(epoch):
return min(1., (epoch - self.n_epochs) /
(opt.begin_decay - self.n_epochs + 1))
self.schedulers = [
optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
for optimizer in [self.optimizerG, self.optimizerD]
]
def move_to(self):
"""Move Tensors to cuda if available"""
self.netG_A = self.netG_A.to(self.device, non_blocking=True)
self.netG_B = self.netG_B.to(self.device, non_blocking=True)
self.netD_A = self.netD_A.to(self.device, non_blocking=True)
self.netD_B = self.netD_B.to(self.device, non_blocking=True)
self.criterionGAN = self.criterionGAN.to(self.device,
non_blocking=True)
self.criterionCycle = self.criterionCycle.to(self.device,
non_blocking=True)
self.criterionIdt = self.criterionIdt.to(self.device,
non_blocking=True)
self.true_label = self.true_label.to(self.device, non_blocking=True)
self.false_label = self.false_label.to(self.device, non_blocking=True)
def backward_G(self):
"""Compute losses and gradients"""
pred_A = self.netD_A(self.fakeB, parallel=opt.parallel)
pred_B = self.netD_B(self.fakeA, parallel=opt.parallel)
self.loss_G_A = self.criterionGAN(
pred_A, self.true_label.repeat(opt.batch_size,
*pred_A[0][0].size()))
self.loss_G_B = self.criterionGAN(
pred_B, self.true_label.repeat(opt.batch_size,
*pred_B[0][0].size()))
self.loss_cycle_A = self.criterionCycle(self.recoA, self.realA)
self.loss_cycle_B = self.criterionCycle(self.recoB, self.realB)
self.loss_idt_A = self.criterionIdt(self.idtA, self.realB)
self.loss_idt_B = self.criterionIdt(self.idtB, self.realA)
self.loss_G = (self.loss_G_A + self.loss_G_B +
self.loss_cycle_A * self.lambdaA +
self.loss_cycle_B * self.lambdaB +
self.loss_idt_A * self.lambdaA * self.lambdaIdt +
self.loss_idt_B * self.lambdaB * self.lambdaIdt)
self.loss_G.backward()
def compute_loss_D_basic(self, netD, real, fake):
"""Compute losses of corresponding discriminator"""
pred_real = netD(real, parallel=opt.parallel)
pred_fake = netD(fake.detach(), parallel=opt.parallel)
loss_D_real = self.criterionGAN(
pred_real,
self.true_label.repeat(opt.batch_size, *pred_real[0][0].size()))
loss_D_fake = self.criterionGAN(
pred_fake,
self.false_label.repeat(opt.batch_size, *pred_fake[0][0].size()))
loss_D = (loss_D_real + loss_D_fake) / 2
return loss_D
def compute_loss_D_A(self):
"""Compute the loss of D_A
Discriminator needs to get an image from the image pool
"""
fake_B = self.fakeB_pool.query(self.fakeB)
self.loss_D_A = self.compute_loss_D_basic(self.netD_A, self.realB,
fake_B)
def compute_loss_D_B(self):
"""Compute the loss of D_B
Discriminator needs to get an image from the image pool
"""
fake_A = self.fakeA_pool.query(self.fakeA)
self.loss_D_B = self.compute_loss_D_basic(self.netD_B, self.realA,
fake_A)
def backward_D(self):
"""Compute the final loss of discriminators and gradients"""
self.compute_loss_D_A()
self.compute_loss_D_B()
self.loss_D = self.loss_D_A + self.loss_D_B
self.loss_D.backward()
def forward(self,
realA: torch.Tensor,
realB: torch.Tensor,
parallel=opt.parallel):
"""Forward images to the net"""
self.realA = realA.to(self.device, non_blocking=True)
self.realB = realB.to(self.device, non_blocking=True)
# X <-------> Y
self.fakeB = self.netG_A(self.realA,
parallel=parallel) # realA --G_A--> fakeB
self.recoA = self.netG_B(self.fakeB,
parallel=parallel) # recoA <--G_B-- fakeB
self.fakeA = self.netG_B(self.realB,
parallel=parallel) # fakeA <--G_B-- realB
self.recoB = self.netG_A(self.fakeA,
parallel=parallel) # fakeA --G_A--> recoB
# to preserve color composition # X Y
self.idtA = self.netG_A(self.realB,
parallel=parallel) # G_B--> idtB realB ----⌍
self.idtB = self.netG_B(
self.realA, parallel=parallel) # ⌎---- realA idtA <--G_A
def backward(self):
"""Optimize the parameters"""
self.optimizerG.zero_grad()
self.backward_G()
self.optimizerG.step()
self.optimizerD.zero_grad()
self.backward_D()
self.optimizerD.step()
def get_current_images(self) -> dict:
"""Returns lately generated images and input images with names"""
return {
'realA': self.realA,
'realB': self.realB,
'fakeA': self.fakeA,
'fakeB': self.fakeB,
'recoA': self.recoA,
'recoB': self.recoB,
'idtA': self.idtA,
'idtB': self.idtB
}
def get_current_losses(self) -> dict:
"""Returns losses of this step with names"""
loss_names = [
'G', 'G_A', 'G_B', 'Cycle_A', 'Cycle_B', 'Idt_A', 'Idt_B', 'D',
'D_A', 'D_B'
]
losses = [
self.loss_G, self.loss_G_A, self.loss_G_B, self.loss_cycle_A,
self.loss_cycle_B, self.loss_idt_A, self.loss_idt_B, self.loss_D,
self.loss_D_A, self.loss_D_B
]
return {loss_name: loss for loss_name, loss in zip(loss_names, losses)}
def update_learning_rate(self):
"""Update the learning rate at the end of each epoch"""
for scheduler in self.schedulers:
scheduler.step()
lr = self.optimizerG.param_groups[0]['lr']
print(f'learning rate = {lr:.7f}')
def save(self, epoch, total_time):
"""
:param epoch: epoch that has just finished. Training will begin from (epoch + 1).
:param total_time:
:return:
"""
savefile = PATH('CycleGAN_ckpt.pth')
torch.save(
{
'epoch': epoch,
'total time': total_time,
'G_A': self.netG_A.state_dict(),
'G_B': self.netG_B.state_dict(),
'D_A': self.netD_A.state_dict(),
'D_B': self.netD_B.state_dict(),
'optimizerG': self.optimizerG.state_dict(),
'optimizerD': self.optimizerD.state_dict()
}, savefile)
def load(self, path='auto') -> tuple:
checkpoint = torch.load(
PATH('CycleGAN_ckpt.pth') if path == 'auto' else path)
self.netG_A.load_state_dict(checkpoint['G_A'])
self.netG_B.load_state_dict(checkpoint['G_B'])
self.netD_A.load_state_dict(checkpoint['D_A'])
self.netD_B.load_state_dict(checkpoint['D_B'])
self.optimizerG.load_state_dict(checkpoint['optimizerG'])
self.optimizerD.load_state_dict(checkpoint['optimizerD'])
return checkpoint['epoch'], checkpoint['total_time']
def train(self):
self.netG_A.train()
self.netG_B.train()
self.netD_A.train()
self.netD_B.train()
def eval(self):
self.netG_A.eval()
self.netG_B.eval()
self.netD_A.eval()
self.netD_B.eval()
transform=train_trans, image_set='train')
train_dl = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=4, drop_last=True)
val_ds = VOCClassification('/path/to/VOC', transform=val_trans, image_set='val')
val_dl = DataLoader(val_ds, batch_size=8, shuffle=True, num_workers=2, drop_last=True)
# Model
if args.arc == 'vgg':
model = vgg19(pretrained=True)
num_ftrs = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_ftrs, train_ds.CLASSES)
model = DataParallelModel(model.cuda())
else:
raise Exception("Architecture {} not found".format(args.arc))
criterion = DataParallelCriterion(nn.BCEWithLogitsLoss().cuda())
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=0.0005)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 20, gamma=0.2)
best_pred = 0
# Load model
if args.resume:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
best_pred = checkpoint['best_pred']
for epoch in range(args.epochs):
class Trainer(CheckpointRunner):
# noinspection PyAttributeOutsideInit
def init_fn(self, shared_model=None, **kwargs):
# Create auxiliary models
self.init_auxiliary()
if shared_model is not None:
self.model = shared_model
else:
self.model = self.init_model()
self.model = DataParallelModel(self.model.cuda(),
device_ids=self.gpus)
# self.model = torch.nn.DataParallel(self.model, device_ids=self.gpus).cuda()
# Setup a joint optimizer for the 2 models
self.optimizer = self.init_optimizer(self.options.optim.name)
self.lr_scheduler = self.init_lr(self.options.optim.lr_scheduler)
# Create loss functions
self.criterion = self.init_loss_functions()
self.criterion = DataParallelCriterion(self.criterion.cuda(),
device_ids=self.gpus)
# Create AverageMeters for losses
self.losses = AverageMeter()
# Evaluators
# self.evaluators = [Evaluator(self.options, self.logger, self.summary_writer, shared_model=self.model)]
self.dataset_size = None
def init_auxiliary(self):
pass
def init_model(self):
raise NotImplementedError("Your model is not found")
def init_loss_functions(self):
raise NotImplementedError("Your loss is not found")
def init_optimizer(self, optim_name):
if optim_name == "adam":
optimizer = torch.optim.Adam(params=list(self.model.parameters()),
lr=self.options.optim.lr,
betas=(self.options.optim.adam_beta1,
0.999),
weight_decay=self.options.optim.wd)
elif optim_name == "sgd":
optimizer = torch.optim.SGD(
params=list(self.model.parameters()),
lr=self.options.optim.lr,
momentum=self.options.optim.sgd_momentum,
weight_decay=self.options.optim.wd)
elif optim_name == "adam_gan":
optimizer_d = torch.optim.Adam(
params=list(self.model.module.D.parameters()),
lr=self.options.optim.lr_d,
betas=(self.options.optim.adam_beta1, 0.999),
weight_decay=0)
optimizer_g = torch.optim.Adam(
params=list(self.model.module.G.parameters()),
lr=self.options.optim.lr_g,
betas=(self.options.optim.adam_beta1, 0.999),
weight_decay=0)
return {"optimizer_d": optimizer_d, "optimizer_g": optimizer_g}
else:
raise NotImplementedError("Your optimizer is not found")
return optimizer
def init_lr(self, lr_scheduler_name):
if lr_scheduler_name == "multistep":
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
self.optimizer, self.options.optim.lr_step,
self.options.optim.lr_factor)
elif lr_scheduler_name == "exp":
lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
self.optimizer, gamma=self.options.optim.lr_gamma)
elif lr_scheduler_name == "multistep_gan":
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
self.optimizer["optimizer_d"], self.options.optim.lr_step,
self.options.optim.lr_factor)
else:
r_scheduler = None
return lr_scheduler
def models_dict(self):
return {'model': self.model}
def optimizers_dict(self):
return {'optimizer': self.optimizer, 'lr_scheduler': self.lr_scheduler}
def train_step(self, input_batch):
# Grab data from the batch, predict with model
out = self.model(input_batch)
# compute loss
loss, loss_summary = self.criterion(out, input_batch)
self.losses.update(loss.detach().cpu().item())
# Do backprop
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Pack output arguments to be used for visualization
return recursive_detach(out), recursive_detach(loss_summary)
def get_dataloader(self):
data_loader = DataLoader(self.dataset,
batch_size=self.options.train.batch_size *
self.options.num_gpus,
num_workers=self.options.num_workers,
pin_memory=self.options.pin_memory,
shuffle=self.options.train.shuffle)
return data_loader
def train(self):
self.logger.info("Start Trainning.")
# Create data loader at very begining
train_data_loader = self.get_dataloader()
self.dataset_size = len(train_data_loader)
# Run training for num_epochs epochs
for epoch in range(self.epoch_count, self.options.train.num_epochs):
self.epoch_count += 1
# Reset loss
self.losses.reset()
# Iterate over all batches in an epoch
for step, batch in enumerate(train_data_loader):
# Send input to GPU
batch = {
k: v.cuda() if isinstance(v, torch.Tensor) else v
for k, v in batch.items()
}
# Run training step
out = self.train_step(batch)
self.step_count += 1
# Tensorboard logging every summary_steps steps
if self.step_count % self.options.train.summary_steps == 0:
self.train_summaries(batch, *out)
# Save checkpoint every checkpoint_steps steps
if self.step_count % self.options.train.checkpoint_steps == 0:
self.dump_checkpoint()
if not self.options.model.name.endswith('gan'):
self.dump_checkpoint()
if self.lr_scheduler is not None:
self.lr_scheduler.step()
def train_summaries(self, input_batch, out_summary, loss_summary):
# Debug info for filenames
self.logger.debug(input_batch["filename"])
# Save results in Tensorboard
self.tensorboard_step(loss_summary)
# Save results to log
self.log_step(loss_summary)
def log_step(self, loss_summary):
self.logger.info(
"Epoch %03d, Step %06d/%06d, Time elapsed %s, Loss %.5f (AvgLoss %.5f)"
% (self.epoch_count, self.step_count,
self.options.train.num_epochs * len(self.dataset) //
(self.options.train.batch_size * self.options.num_gpus),
self.time_elapsed, self.losses.val, self.losses.avg))
def tensorboard_step(self, loss_summary):
for k, v in loss_summary.items():
self.summary_writer.add_scalar(k, v, self.step_count)
def init_with_pretrained_backbone(self):
checkpoint_file = os.path.abspath(
self.options.train.backbone_pretrained_model)
pretrained_dict = torch.load(checkpoint_file)
self.model.module.load_state_dict(pretrained_dict, strict=False)
self.logger.info("Init with pre-trained backbone from %s." %
checkpoint_file)
def test(self):
self.model.eval()
for evaluator in self.evaluators:
evaluator.evaluate()
self.model.train()
def main(args):
# initialization
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
writer = SummaryWriter(log_dir=os.path.join(args.log_dir, args.method))
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.enabled = True
cudnn.benchmark = True
adj_matrix = torch.tensor(
[[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0],
[0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1],
[0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0],
[0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0],
[1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1],
[0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1],
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0]],
requires_grad=False)
upper_part_list = [1, 2, 3, 4, 5, 6, 7, 11, 13, 14, 15]
lower_part_list = [8, 9, 10, 12, 16, 17, 18, 19]
weight = torch.FloatTensor([
0.7602572, 0.94236198, 0.85644457, 1.04346266, 1.10627293, 0.80980162,
0.95168713, 0.8403769, 1.05798412, 0.85746254, 1.01274366, 1.05854692,
1.03430773, 0.84867818, 0.88027721, 0.87580925, 0.98747462, 0.9876475,
1.00016535, 1.00108882
])
# conduct seg network
seg_model = get_model(num_classes=args.num_classes,
adj_matrix=adj_matrix,
upper_part_list=upper_part_list,
lower_part_list=lower_part_list)
saved_state_dict = torch.load(args.restore_from)
new_params = seg_model.state_dict().copy()
# if args.init:
# for i in saved_state_dict:
# i_parts = i.split('.')
# if not i_parts[0] == 'fc':
# new_params['encoder.' + '.'.join(i_parts[:])] = saved_state_dict[i]
# #new_params[i_parts[:]] = saved_state_dict[i]
# seg_model.load_state_dict(new_params)
# print('loading params w/o fc')
# else:
# seg_model.load_state_dict(saved_state_dict)
# print('loading params all')
for i in saved_state_dict:
i_parts = i.split('.')
if not i_parts[0] == 'fc':
new_params['encoder.' + '.'.join(i_parts[:])] = saved_state_dict[i]
seg_model.load_state_dict(new_params)
print('loading params w/o fc')
# seg_model.load_state_dict(saved_state_dict)
# print('loading params all')
model = DataParallelModel(seg_model)
model.float()
model.cuda()
# define dataloader
train_loader = data.DataLoader(DatasetGenerator(root=args.root,
list_path=args.lst,
crop_size=args.crop_size,
training=True),
batch_size=args.batch_size,
shuffle=True,
num_workers=4,
pin_memory=True)
val_loader = data.DataLoader(DatasetGenerator(root=args.val_root,
list_path=args.val_lst,
crop_size=args.crop_size,
training=False),
batch_size=args.batch_size,
shuffle=False,
num_workers=4,
pin_memory=True)
# define criterion & optimizer
criterion = ABRLovaszLoss(adj_matrix=adj_matrix,
ignore_index=args.ignore_label,
only_present=True,
upper_part_list=upper_part_list,
lower_part_list=lower_part_list,
cls_p=args.num_classes,
cls_h=args.hbody_cls,
cls_f=args.fbody_cls,
weight=weight)
criterion = DataParallelCriterion(criterion).cuda()
optimizer = optim.SGD(
[{
'params': filter(lambda p: p.requires_grad,
seg_model.parameters()),
'lr': args.learning_rate
}],
lr=args.learning_rate,
momentum=0.9,
weight_decay=5e-4)
# key points
best_val_mIoU = 0
best_val_pixAcc = 0
start = time.time()
for epoch in range(0, args.epochs):
print('\n{} | {}'.format(epoch, args.epochs - 1))
# training
_ = train(model, train_loader, epoch, criterion, optimizer, writer)
# validation
if epoch % 10 == 0 or epoch > args.epochs - 5:
val_pixacc, val_miou = validation(model, val_loader, epoch, writer)
# save model
if val_pixacc > best_val_pixAcc:
best_val_pixAcc = val_pixacc
if val_miou > best_val_mIoU:
best_val_mIoU = val_miou
model_dir = os.path.join(args.snapshot_dir,
args.method + '_miou.pth')
torch.save(seg_model.state_dict(), model_dir)
print('Model saved to %s' % model_dir)
os.rename(
model_dir,
os.path.join(args.snapshot_dir,
args.method + '_miou' + str(best_val_mIoU) + '.pth'))
print('Complete using', time.time() - start, 'seconds')
print('Best pixAcc: {} | Best mIoU: {}'.format(best_val_pixAcc,
best_val_mIoU))
