def main():
import argparse
parser = argparse.ArgumentParser(
description="imsitu VSRL. Training, evaluation and prediction.")
parser.add_argument("--gpuid",
default=-1,
help="put GPU id > -1 in GPU mode",
type=int)
parser.add_argument('--output_dir',
type=str,
default='./trained_models',
help='Location to output the model')
parser.add_argument('--resume_training',
action='store_true',
help='Resume training from the model [resume_model]')
parser.add_argument('--resume_model',
type=str,
default='',
help='The model we resume')
parser.add_argument('--evaluate',
action='store_true',
help='Only use the testing mode')
parser.add_argument('--evaluate_rare',
action='store_true',
help='Only use the testing mode')
parser.add_argument('--test',
action='store_true',
help='Only use the testing mode')
parser.add_argument('--dataset_folder',
type=str,
default='./imSitu',
help='Location of annotations')
parser.add_argument('--imgset_dir',
type=str,
default='./resized_256',
help='Location of original images')
parser.add_argument('--train_file',
default="train_freq2000.json",
type=str,
help='trainfile name')
parser.add_argument('--dev_file',
default="dev_freq2000.json",
type=str,
help='dev file name')
parser.add_argument('--test_file',
default="test_freq2000.json",
type=str,
help='test file name')
parser.add_argument('--model_saving_name',
type=str,
help='saving name of the outpul model')
parser.add_argument('--epochs', type=int, default=500)
parser.add_argument('--model', type=str, default='vgg_caq_joint')
parser.add_argument('--batch_size', type=int, default=64)
parser.add_argument('--seed', type=int, default=1111, help='random seed')
parser.add_argument('--clip_norm', type=float, default=0.25)
parser.add_argument('--num_workers', type=int, default=3)
parser.add_argument('--vgg_verb_model',
type=str,
default='',
help='Pretrained vgg verb model')
parser.add_argument('--tda_verb_model',
type=str,
default='',
help='Pretrained topdown verb model')
parser.add_argument('--caq_model',
type=str,
default='',
help='Pretrained CAQ model')
args = parser.parse_args()
n_epoch = args.epochs
batch_size = args.batch_size
clip_norm = args.clip_norm
n_worker = args.num_workers
dataset_folder = args.dataset_folder
imgset_folder = args.imgset_dir
train_set = json.load(open(dataset_folder + '/' + args.train_file))
encoder = imsitu_encoder.imsitu_encoder(train_set)
train_set = imsitu_loader.imsitu_loader(imgset_folder, train_set, encoder,
'train', encoder.train_transform)
constructor = 'build_vgg_verb_classifier'
vgg_verb_model = getattr(vgg_verb_classifier,
constructor)(len(encoder.verb_list))
constructor = 'build_top_down_baseline'
role_module = getattr(top_down_baseline_addemb,
constructor)(encoder.get_num_roles(),
encoder.get_num_verbs(),
encoder.get_num_labels(), encoder)
constructor = 'build_top_down_verb'
tda_verb_model = getattr(top_down_verb,
constructor)(encoder.get_num_labels(),
encoder.get_num_verbs(), role_module)
constructor = 'build_top_down_baseline'
tda_role_module = getattr(top_down_baseline,
constructor)(encoder.get_num_roles(),
encoder.get_num_verbs(),
encoder.get_num_labels(), encoder)
constructor = 'build_top_down_query_context'
caq_model = getattr(top_down_query_context,
constructor)(encoder.get_num_roles(),
encoder.get_num_verbs(),
encoder.get_num_labels(), encoder,
tda_role_module)
constructor = 'build_%s' % args.model
model = getattr(revgg_caq_joint_eval,
constructor)(vgg_verb_model, tda_verb_model, caq_model)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=n_worker)
dev_set = json.load(open(dataset_folder + '/' + args.dev_file))
dev_set = imsitu_loader.imsitu_loader(imgset_folder, dev_set, encoder,
'val', encoder.dev_transform)
dev_loader = torch.utils.data.DataLoader(dev_set,
batch_size=batch_size,
shuffle=True,
num_workers=n_worker)
test_set = json.load(open(dataset_folder + '/' + args.test_file))
test_set = imsitu_loader.imsitu_loader(imgset_folder, test_set, encoder,
'test', encoder.dev_transform)
test_loader = torch.utils.data.DataLoader(test_set,
batch_size=batch_size,
shuffle=True,
num_workers=n_worker)
if not os.path.exists(args.output_dir):
os.mkdir(args.output_dir)
torch.manual_seed(args.seed)
if args.gpuid >= 0:
model.cuda()
torch.cuda.manual_seed(args.seed)
torch.backends.cudnn.benchmark = True
#load models
utils.load_net(args.vgg_verb_model, [model.vgg_model])
print('successfully loaded vgg_verb_model!')
utils.load_net(args.tda_verb_model, [model.tda_model])
print('successfully loaded tda_verb_model!')
utils.load_net(args.caq_model, [model.caq_model])
print('successfully loaded caq_model!')
if args.evaluate:
top1, top5, val_loss = eval(model,
dev_loader,
encoder,
args.gpuid,
write_to_file=True)
top1_avg = top1.get_average_results()
top5_avg = top5.get_average_results()
avg_score = top1_avg["verb"] + top1_avg["value"] + top1_avg["value-all"] + top5_avg["verb"] + \
top5_avg["value"] + top5_avg["value-all"] + top5_avg["value*"] + top5_avg["value-all*"]
avg_score /= 8
print('Dev average :{:.2f} {} {}'.format(
avg_score * 100, utils.format_dict(top1_avg, '{:.2f}', '1-'),
utils.format_dict(top5_avg, '{:.2f}', '5-')))
elif args.test:
top1, top5, val_loss = eval(model,
test_loader,
encoder,
args.gpuid,
write_to_file=True)
top1_avg = top1.get_average_results()
top5_avg = top5.get_average_results()
avg_score = top1_avg["verb"] + top1_avg["value"] + top1_avg["value-all"] + top5_avg["verb"] + \
top5_avg["value"] + top5_avg["value-all"] + top5_avg["value*"] + top5_avg["value-all*"]
avg_score /= 8
print('Test average :{:.2f} {} {}'.format(
avg_score * 100, utils.format_dict(top1_avg, '{:.2f}', '1-'),
utils.format_dict(top5_avg, '{:.2f}', '5-')))
def train(model,
train_loader,
dev_loader,
optimizer,
scheduler,
max_epoch,
model_dir,
encoder,
gpu_mode,
clip_norm,
model_name,
model_saving_name,
eval_frequency=4000):
model.train()
train_loss = 0
total_steps = 0
print_freq = 400
dev_score_list = []
if gpu_mode >= 0:
ngpus = 2
device_array = [i for i in range(0, ngpus)]
pmodel = torch.nn.DataParallel(model, device_ids=device_array)
else:
pmodel = model
top1 = imsitu_scorer.imsitu_scorer(encoder, 1, 3)
top5 = imsitu_scorer.imsitu_scorer(encoder, 5, 3)
for epoch in range(max_epoch):
for i, (_, img, verb, labels) in enumerate(train_loader):
total_steps += 1
if gpu_mode >= 0:
img = torch.autograd.Variable(img.cuda())
verb = torch.autograd.Variable(verb.cuda())
labels = torch.autograd.Variable(labels.cuda())
else:
img = torch.autograd.Variable(img)
verb = torch.autograd.Variable(verb)
labels = torch.autograd.Variable(labels)
role_predict = pmodel(img, verb)
loss = model.calculate_loss(verb, role_predict, labels)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
optimizer.step()
optimizer.zero_grad()
train_loss += loss.item()
top1.add_point_noun(verb, role_predict, labels)
top5.add_point_noun(verb, role_predict, labels)
if total_steps % print_freq == 0:
top1_a = top1.get_average_results_nouns()
top5_a = top5.get_average_results_nouns()
print("{},{},{}, {} , {}, loss = {:.2f}, avg loss = {:.2f}".
format(total_steps - 1, epoch, i,
utils.format_dict(top1_a, "{:.2f}", "1-"),
utils.format_dict(top5_a, "{:.2f}", "5-"),
loss.item(), train_loss /
((total_steps - 1) % eval_frequency)))
if total_steps % eval_frequency == 0:
top1, top5, val_loss = eval(model, dev_loader, encoder,
gpu_mode)
model.train()
top1_avg = top1.get_average_results_nouns()
top5_avg = top5.get_average_results_nouns()
avg_score = top1_avg["verb"] + top1_avg["value"] + top1_avg["value-all"] + top5_avg["verb"] + \
top5_avg["value"] + top5_avg["value-all"] + top5_avg["value*"] + top5_avg["value-all*"]
avg_score /= 8
print('Dev {} average :{:.2f} {} {}'.format(
total_steps - 1, avg_score * 100,
utils.format_dict(top1_avg, '{:.2f}', '1-'),
utils.format_dict(top5_avg, '{:.2f}', '5-')))
dev_score_list.append(avg_score)
max_score = max(dev_score_list)
if max_score == dev_score_list[-1]:
torch.save(
model.state_dict(), model_dir +
"/{}_{}.model".format(model_name, model_saving_name))
print('New best model saved! {0}'.format(max_score))
print('current train loss', train_loss)
train_loss = 0
top1 = imsitu_scorer.imsitu_scorer(encoder, 1, 3)
top5 = imsitu_scorer.imsitu_scorer(encoder, 5, 3)
del role_predict, loss, img, verb, labels
print('Epoch ', epoch, ' completed!')
scheduler.step()
def printdict(d):
return format_dict(d)
def main():
import argparse
parser = argparse.ArgumentParser(
description="imsitu VSRL. Training, evaluation and prediction.")
parser.add_argument("--gpuid",
default=-1,
help="put GPU id > -1 in GPU mode",
type=int)
parser.add_argument('--output_dir',
type=str,
default='./trained_models',
help='Location to output the model')
parser.add_argument('--resume_training',
action='store_true',
help='Resume training from the model [resume_model]')
parser.add_argument('--resume_model',
type=str,
default='',
help='The model we resume')
parser.add_argument('--evaluate',
action='store_true',
help='Only use the testing mode')
parser.add_argument('--evaluate_visualize',
action='store_true',
help='Only use the testing mode to visualize ')
parser.add_argument('--evaluate_rare',
action='store_true',
help='Only use the testing mode')
parser.add_argument('--test',
action='store_true',
help='Only use the testing mode')
parser.add_argument('--dataset_folder',
type=str,
default='./imSitu',
help='Location of annotations')
parser.add_argument('--imgset_dir',
type=str,
default='./resized_256',
help='Location of original images')
parser.add_argument('--train_file',
default="train_freq2000.json",
type=str,
help='trainfile name')
parser.add_argument('--dev_file',
default="dev_freq2000.json",
type=str,
help='dev file name')
parser.add_argument('--test_file',
default="test_freq2000.json",
type=str,
help='test file name')
parser.add_argument('--model_saving_name',
type=str,
help='saving name of the outpul model')
parser.add_argument('--epochs', type=int, default=500)
parser.add_argument('--model', type=str, default='top_down_baseline')
parser.add_argument('--batch_size', type=int, default=64)
parser.add_argument('--seed', type=int, default=1111, help='random seed')
parser.add_argument('--clip_norm', type=float, default=0.25)
parser.add_argument('--num_workers', type=int, default=3)
args = parser.parse_args()
n_epoch = args.epochs
batch_size = args.batch_size
clip_norm = args.clip_norm
n_worker = args.num_workers
dataset_folder = args.dataset_folder
imgset_folder = args.imgset_dir
train_set = json.load(open(dataset_folder + '/' + args.train_file))
encoder = imsitu_encoder.imsitu_encoder(train_set)
train_set = imsitu_loader.imsitu_loader(imgset_folder, train_set, encoder,
'train', encoder.train_transform)
constructor = 'build_%s' % args.model
model = getattr(top_down_baseline, constructor)(encoder.get_num_roles(),
encoder.get_num_verbs(),
encoder.get_num_labels(),
encoder)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=n_worker)
dev_set = json.load(open(dataset_folder + '/' + args.dev_file))
dev_set = imsitu_loader.imsitu_loader(imgset_folder, dev_set, encoder,
'val', encoder.dev_transform)
dev_loader = torch.utils.data.DataLoader(dev_set,
batch_size=batch_size,
shuffle=True,
num_workers=n_worker)
test_set = json.load(open(dataset_folder + '/' + args.test_file))
test_set = imsitu_loader.imsitu_loader(imgset_folder, test_set, encoder,
'test', encoder.dev_transform)
test_loader = torch.utils.data.DataLoader(test_set,
batch_size=batch_size,
shuffle=True,
num_workers=n_worker)
if not os.path.exists(args.output_dir):
os.mkdir(args.output_dir)
torch.manual_seed(args.seed)
if args.gpuid >= 0:
model.cuda()
torch.cuda.manual_seed(args.seed)
torch.backends.cudnn.benchmark = True
if args.resume_training:
print('Resume training from: {}'.format(args.resume_model))
args.train_all = True
if len(args.resume_model) == 0:
raise Exception('[pretrained module] not specified')
utils.load_net(args.resume_model, [model])
optimizer = torch.optim.Adamax(model.parameters(), lr=1e-3)
model_name = 'resume_all'
else:
print('Training from the scratch.')
model_name = 'train_full'
utils.set_trainable(model, True)
optimizer = torch.optim.Adamax(
[{
'params': model.convnet.parameters(),
'lr': 5e-5
}, {
'params': model.role_emb.parameters()
}, {
'params': model.verb_emb.parameters()
}, {
'params': model.query_composer.parameters()
}, {
'params': model.v_att.parameters()
}, {
'params': model.q_net.parameters()
}, {
'params': model.v_net.parameters()
}, {
'params': model.classifier.parameters()
}],
lr=1e-3)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)
if args.evaluate:
top1, top5, val_loss = eval(model, dev_loader, encoder, args.gpuid)
top1_avg = top1.get_average_results_nouns()
top5_avg = top5.get_average_results_nouns()
avg_score = top1_avg["verb"] + top1_avg["value"] + top1_avg["value-all"] + top5_avg["verb"] + \
top5_avg["value"] + top5_avg["value-all"] + top5_avg["value*"] + top5_avg["value-all*"]
avg_score /= 8
print('Dev average :{:.2f} {} {}'.format(
avg_score * 100, utils.format_dict(top1_avg, '{:.2f}', '1-'),
utils.format_dict(top5_avg, '{:.2f}', '5-')))
elif args.test:
top1, top5, val_loss = eval(model, test_loader, encoder, args.gpuid)
top1_avg = top1.get_average_results_nouns()
top5_avg = top5.get_average_results_nouns()
avg_score = top1_avg["verb"] + top1_avg["value"] + top1_avg["value-all"] + top5_avg["verb"] + \
top5_avg["value"] + top5_avg["value-all"] + top5_avg["value*"] + top5_avg["value-all*"]
avg_score /= 8
print('Test average :{:.2f} {} {}'.format(
avg_score * 100, utils.format_dict(top1_avg, '{:.2f}', '1-'),
utils.format_dict(top5_avg, '{:.2f}', '5-')))
else:
print('Model training started!')
train(
model,
train_loader,
dev_loader,
optimizer,
scheduler,
n_epoch,
args.output_dir,
encoder,
args.gpuid,
clip_norm,
model_name,
args.model_saving_name,
)
def eval(model, loader, encoder, logging=False):
model.eval()
verbloss = 0
nounsloss = 0
gtloss = 0
top1 = imsitu_scorer.imsitu_scorer(encoder, 1, 3)
top5 = imsitu_scorer.imsitu_scorer(encoder, 5, 3)
with torch.no_grad():
for __, (_, img, verb, nouns) in enumerate(loader):
if torch.cuda.is_available():
img = img.cuda()
verb = verb.cuda()
nouns = nouns.cuda()
with autocast(): # automix precision stuff
pred_verb, pred_nouns, pred_gt_nouns = model(img, verb)
top1.add_point_both(pred_verb, verb, pred_nouns, nouns,
pred_gt_nouns)
top5.add_point_both(pred_verb, verb, pred_nouns, nouns,
pred_gt_nouns)
if torch.cuda.is_available():
vl = model.module.verb_loss(pred_verb, verb)
nl = model.module.nouns_loss(pred_nouns, nouns)
gtl = model.module.nouns_loss(pred_gt_nouns, nouns)
else:
vl = model.verb_loss(pred_verb, verb)
nl = model.nouns_loss(pred_nouns, nouns)
gtl = model.nouns_loss(pred_gt_nouns, nouns)
verbloss += vl.item()
nounsloss += nl.item()
gtloss += gtl.item()
verbloss /= len(loader)
nounsloss /= len(loader)
gtloss /= len(loader)
val_losses = {
'verb_loss': verbloss,
'nouns_loss': nounsloss,
'gt_loss': gtloss
}
#print scores
avg_score = 0
if logging is True:
top1_a = top1.get_average_results_both()
top5_a = top5.get_average_results_both()
avg_score = top1_a['verb'] + top1_a['value'] + top1_a['value-all'] + \
top5_a['verb'] + top5_a['value'] + top5_a['value-all'] + \
top1_a['gt-value'] + top1_a['gt-value-all']
avg_score /= 8
avg_score = avg_score * 100
print('val losses = [v: {:.2f}, n: {:.2f}, gt: {:.2f}]'.format(
val_losses['verb_loss'], val_losses['nouns_loss'],
val_losses['gt_loss']))
gt = {key: top1_a[key] for key in ['gt-value', 'gt-value-all']}
one_val = {key: top1_a[key] for key in ['verb', 'value', 'value-all']}
print('{}\n{}\n{}, mean = {:.2f}\n'.format(
utils.format_dict(one_val, '{:.2f}', '1-'),
utils.format_dict(top5_a, '{:.2f}', '5-'),
utils.format_dict(gt, '{:.2f}', ''), avg_score))
return top1, top5, val_losses, avg_score
def train(model,
train_loader,
dev_loader,
optimizer,
max_epoch,
encoder,
model_saving_name,
folder,
checkpoint=None):
model.train()
avg_scores = []
verb_losses = []
nouns_losses = []
val_avg_scores = []
val_verb_losses = []
val_nouns_losses = []
epoch = 0
# if checkpoint resume stuffs
if checkpoint is not None:
epoch = checkpoint['epoch']
avg_scores = checkpoint['avg_scores']
verb_losses = checkpoint['verb_losses']
nouns_losses = checkpoint['nouns_losses']
val_avg_scores = checkpoint['val_avg_scores']
val_verb_losses = checkpoint['val_verb_losses']
val_nouns_losses = checkpoint['val_nouns_losses']
if torch.cuda.is_available():
model.module.load_state_dict(checkpoint['model_state_dict'])
else:
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
#mix precision stuff
scaler = GradScaler()
for e in range(epoch, max_epoch):
verb_loss_accum = 0
nouns_loss_accum = 0
gt_nouns_loss_accum = 0
print('Epoch-{}, lr: {:.4f}'.format(e,
optimizer.param_groups[0]['lr']))
top1 = imsitu_scorer.imsitu_scorer(encoder, 1, 3)
top5 = imsitu_scorer.imsitu_scorer(encoder, 5, 3)
for i, (_, img, verb, nouns) in enumerate(train_loader):
if torch.cuda.is_available():
img = img.cuda()
verb = verb.cuda()
nouns = nouns.cuda()
optimizer.zero_grad()
with autocast(): #mix precision stuff
pred_verb, pred_nouns, pred_gt_nouns = model(img, verb)
#predict and calculate lossess
if torch.cuda.is_available():
verb_loss = model.module.verb_loss(pred_verb, verb)
nouns_loss = model.module.nouns_loss(pred_nouns, nouns)
gt_nouns_loss = model.module.nouns_loss(
pred_gt_nouns, nouns)
else:
verb_loss = model.verb_loss(pred_verb, verb)
nouns_loss = model.nouns_loss(pred_nouns, nouns)
gt_nouns_loss = model.nouns_loss(pred_gt_nouns, nouns)
loss = verb_loss + nouns_loss
# backpropagate errors and stuffs
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1)
scaler.step(optimizer)
scaler.update()
top1.add_point_both(pred_verb, verb, pred_nouns, nouns,
pred_gt_nouns)
top5.add_point_both(pred_verb, verb, pred_nouns, nouns,
pred_gt_nouns)
verb_loss_accum += verb_loss.item()
nouns_loss_accum += nouns_loss.item()
gt_nouns_loss_accum += gt_nouns_loss.item()
#fin epoch
#epoch accuracy mean
top1_a = top1.get_average_results_both()
top5_a = top5.get_average_results_both()
avg_score = top1_a['verb'] + top1_a['value'] + top1_a['value-all'] + \
top5_a['verb'] + top5_a['value'] + top5_a['value-all'] + \
top1_a['gt-value'] + top1_a['gt-value-all']
avg_score /= 8
avg_score *= 100
avg_scores.append(avg_score)
#epoch lossess
verb_loss_mean = verb_loss_accum / len(train_loader)
nouns_loss_mean = nouns_loss_accum / len(train_loader)
gt_nouns_loss_mean = gt_nouns_loss_accum / len(train_loader)
verb_losses.append(verb_loss_mean)
nouns_losses.append(nouns_loss_mean)
#print stuffs
print('training losses = [v: {:.2f}, n: {:.2f}, gt: {:.2f}]'.format(
verb_loss_mean, nouns_loss_mean, gt_nouns_loss_mean))
gt = {key: top1_a[key] for key in ['gt-value', 'gt-value-all']}
one_val = {key: top1_a[key] for key in ['verb', 'value', 'value-all']}
print('{}\n{}\n{}, mean = {:.2f}\n{}'.format(
utils.format_dict(one_val, '{:.2f}', '1-'),
utils.format_dict(top5_a, '{:.2f}', '5-'),
utils.format_dict(gt, '{:.2f}', ''), avg_score, '-' * 50))
# evaluating
top1, top5, val_losses, val_avg_score = eval(model,
dev_loader,
encoder,
logging=True)
model.train()
#val mean scores
val_avg_scores.append(val_avg_score)
val_verb_losses.append(val_losses['verb_loss'])
val_nouns_losses.append(val_losses['nouns_loss'])
plt.plot(verb_losses, label='verb losses')
plt.plot(nouns_losses, label='nouns losses')
plt.plot(avg_scores, label='accuracy mean')
plt.plot(val_verb_losses, '-.', label='val verb losses')
plt.plot(val_nouns_losses, '-.', label='val nouns losses')
plt.plot(val_avg_scores, '-.', label='val accuracy mean')
plt.grid()
plt.legend()
plt.savefig(pjoin(folder, model_saving_name + '.png'))
plt.clf()
#always saving but no need if it's not the best score
checkpoint = {
'epoch': e + 1,
'avg_scores': avg_scores,
'verb_losses': verb_losses,
'nouns_losses': nouns_losses,
'val_avg_scores': val_avg_scores,
'val_verb_losses': val_verb_losses,
'val_nouns_losses': val_nouns_losses,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict()
}
if torch.cuda.is_available():
checkpoint.update({'model_state_dict': model.module.state_dict()})
torch.save(checkpoint, pjoin(folder, model_saving_name))
def update_network(self):
# initial configuration
stop = False
update_iters = 0
self.iter('S')
self.iter('T')
while not stop:
# update learning rate
self.update_lr(self.optimizer, self.base_lr)
# set the status of network
self.net.train()
self.net.zero_grad()
if self.adv_train:
self.update_lr(self.optimizer_D, self.base_lr_D)
self.net_D.train()
self.net_D.zero_grad()
loss = 0
for k in range(self.iter_size):
sample_S = self.get_training_samples('S')
data_S, gt_S = sample_S['Img'], sample_S['Label']
data_S, gt_S = gen_utils.to_cuda(data_S), gen_utils.to_cuda(
gt_S)
sample_T = self.get_training_samples('T')
data_T, gt_T = sample_T['Img'], sample_T['Label']
data_T, gt_T = gen_utils.to_cuda(data_T), gen_utils.to_cuda(
gt_T)
loss_dict, out_dict = eval('self.%s' % cfg.TRAIN.METHOD)(
data_S, gt_S, data_T, gt_T)
loss = loss_dict['total'] / self.iter_size
preds_S, preds_T = out_dict['preds_S'], out_dict['preds_T']
if self.adv_train:
# G step:
probs_S, probs_T = F.softmax(preds_S,
dim=1), F.softmax(preds_T,
dim=1)
for param in self.net_D.parameters():
param.requires_grad = False
loss_GD = self.G_step(probs_S, probs_T) / self.iter_size
loss += cfg.TRAIN.ADV_W * loss_GD
loss_dict['G_loss'] = loss_GD
loss.backward()
if self.adv_train:
# D step:
for param in self.net_D.parameters():
param.requires_grad = True
loss_D = self.D_step(probs_S, probs_T) / self.iter_size
loss_dict['D_loss'] = loss_D
loss_D.backward()
# update the network
self.optimizer.step()
if self.adv_train:
# update the discriminator
self.optimizer_D.step()
if cfg.TRAIN.LOGGING and (update_iters+1) % \
(max(1, self.iters_per_epoch // cfg.TRAIN.NUM_LOGGING_PER_EPOCH)) == 0:
preds = out_dict['preds_S']
accu = 100.0 * gen_utils.model_eval(
torch.max(preds, dim=1).indices, gt_S, 'accuracy',
preds.size(1), cfg.DATASET.IGNORE_LABEL).item()
miou = 100.0 * gen_utils.model_eval(
torch.max(preds, dim=1).indices, gt_S, 'mIoU',
preds.size(1), cfg.DATASET.IGNORE_LABEL)[0].item()
cur_loss = loss_dict
eval_res = {'accu': accu, 'miou': miou}
self.logging(cur_loss, eval_res)
if cfg.TRAIN.TEST_INTERVAL > 0 and \
(self.iters+1) % int(cfg.TRAIN.TEST_INTERVAL * self.iters_per_epoch) == 0:
with torch.no_grad():
accu, miou = self.test()
print('Test at (epoch %d, iter %d) with %s.' %
(self.epochs, self.iters,
gen_utils.format_dict({
'accu': accu,
'miou': miou
})))
if not self.distributed or gen_utils.is_main_process():
if cfg.TRAIN.SAVE_CKPT_INTERVAL > 0 and \
(self.iters+1) % int(cfg.TRAIN.SAVE_CKPT_INTERVAL * self.iters_per_epoch) == 0:
self.save_ckpt()
update_iters += 1
self.iters += 1
# update stop condition
if update_iters >= self.iters_per_epoch:
stop = True
else:
stop = False
def logging(self, loss, res):
print('[epoch: %d, iter: %d]: ' % (self.epochs, self.iters))
info_str = gen_utils.format_dict(loss) + '; ' + gen_utils.format_dict(
res)
print(info_str)
