def train_epoch(epoch, steps):
model.train()
loss_avg = 0.0
acc_avg = 0.0
counter = 0
train_loader_iter = iter(train_loader)
for _ in trange(len(train_loader_iter)):
steps += 1
counter += 1
images, targets = next(train_loader_iter)
images = images.to(device)
targets = targets.to(device)
model_output = model(images)
loss = criterion(model_output, targets)
loss_avg += loss.item()
acc = criteria.calculate_acc(model_output, targets)
acc_avg += acc.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
writer.add_scalar("Train Loss", loss.item(), steps)
writer.add_scalar("Train Acc", acc.item(), steps)
print "Epoch {}, Total Iter: {}, Train Avg Loss: {:.6f}".format(
epoch, counter, loss_avg / float(counter))
return steps
def validate_epoch(epoch, steps):
model.eval()
loss_avg = 0.0
acc_avg = 0.0
xe_avg = 0.0
counter = 0
val_loader_iter = iter(val_loader)
for _ in trange(len(val_loader_iter)):
counter += 1
images, targets, _ = next(val_loader_iter)
images = images.to(device)
targets = targets.to(device)
model_output = model(images)
scene_prob, scene_logprob = env.prepare(model_output)
action, action_logprob, _ = env.action(scene_logprob, sample=False)
pred = env.step(scene_prob, action)
loss, scores, xe_loss_item = env.loss(action[0], action_logprob,
pred, scene_prob, targets,
error_function)
loss_avg += loss.item()
acc = criteria.calculate_acc(scores, targets)
acc_avg += acc.item()
xe_avg += xe_loss_item
print("Epoch {}, Valid Avg XE: {:.6f}, Valid Avg Acc: {:.4f}".format(
epoch, xe_avg / float(counter), acc_avg / float(counter)))
def train_epoch(epoch, steps):
model.train()
xe_avg = 0.0
counter = 0
train_loader_iter = iter(train_loader)
for _ in trange(len(train_loader_iter)):
steps += 1
counter += 1
images, targets, all_action_rule = next(train_loader_iter)
images = images.to(device)
targets = targets.to(device)
all_action_rule_device = []
for action_rule in all_action_rule:
action_rule = action_rule.to(device)
all_action_rule_device.append(action_rule)
model_output = model(images)
scene_prob, scene_logprob = env.prepare(model_output)
action, action_logprob, all_action_prob = env.action(scene_logprob)
pred = env.step(scene_prob, action)
loss, scores, xe_loss_item = env.loss(action[0], action_logprob,
pred, scene_prob, targets,
error_function)
aux_loss = criteria.aux_loss(all_action_prob,
all_action_rule_device)
final_loss = loss + args.aux * aux_loss
acc = criteria.calculate_acc(scores, targets)
xe_avg += xe_loss_item
optimizer.zero_grad()
final_loss.backward()
optimizer.step()
print("Epoch {}, Total Iter: {}, Train Avg XE: {:.6f}".format(
epoch, counter, xe_avg / float(counter)))
return steps
def test_epoch(epoch, steps):
model.eval()
loss_avg = 0.0
acc_avg = 0.0
counter = 0
test_loader_iter = iter(test_loader)
for _ in trange(len(test_loader_iter)):
counter += 1
images, targets = next(test_loader_iter)
images = images.to(device)
targets = targets.to(device)
model_output = model(images)
loss = criterion(model_output, targets)
loss_avg += loss.item()
acc = criteria.calculate_acc(model_output, targets)
acc_avg += acc.item()
writer.add_scalar("Test Avg Loss", loss_avg / float(counter), steps)
writer.add_scalar("Test Avg Acc", acc_avg / float(counter), steps)
print "Epoch {}, Test Avg Loss: {:.6f}, Test Avg Acc: {:.4f}".format(
epoch, loss_avg / float(counter), acc_avg / float(counter))
def evaluate(self, subset):
self.model.eval()
counter = 0
loss_avg = 0.0
loss_meta_avg = 0.0
acc_avg = 0.0
acc_multihead_avg = [0.0] * 3
acc_regime = init_acc_regime(self.args.dataset)
if subset == 'train':
loader = self.trainloader
elif subset == 'val':
loader = self.validloader
else:
loader = self.testloader
for batch_data in tqdm(loader, subset):
counter += 1
image, target, meta_target, structure_encoded, data_file = batch_data
image = renormalize(image)
image = image.cuda()
target = target.cuda()
if self.use_meta:
meta_target = meta_target.cuda()
with torch.no_grad():
model_outputs = self.model(image)
if len(model_outputs) == 3:
model_output, meta_pred, model_output_heads = model_outputs
else:
model_output, meta_pred = model_outputs
model_output_heads = None
loss = self.criterion(model_output, target)
loss_avg += loss.item()
acc = criteria.calculate_acc(model_output, target)
acc_avg += acc.item()
if self.use_meta:
if len(meta_pred.shape) > 2:
meta_pred = meta_pred[range(meta_pred.shape[0]), target]
loss_meta = self.criterion_meta(meta_pred, meta_target)
loss_meta_avg += loss_meta.item()
if self.args.multihead:
acc_multihead = [
criteria.calculate_acc(x, target)
for x in model_output_heads
]
for i, x in enumerate(acc_multihead):
acc_multihead_avg[i] += x.item()
if acc_regime is not None:
update_acc_regime(self.args.dataset, acc_regime, model_output,
target, structure_encoded, data_file)
if counter > 0:
if self.use_meta:
print(
"{} - Avg Loss: {:.6f} META {:.6f}, Test Avg Acc: {:.4f}".
format(subset, loss_avg / float(counter),
loss_meta_avg / float(counter),
acc_avg / float(counter)))
else:
print("{} - Avg Loss: {:.6f}, Test Avg Acc: {:.4f}".format(
subset, loss_avg / float(counter),
acc_avg / float(counter)))
if self.args.multihead:
print(
f'Multihead: {[x / float(counter) for x in acc_multihead_avg]}'
)
if acc_regime is not None:
for key in acc_regime.keys():
if acc_regime[key] is not None:
if acc_regime[key][1] > 0:
acc_regime[key] = float(
acc_regime[key][0]) / acc_regime[key][1] * 100
else:
acc_regime[key] = None
return loss_avg / float(counter), acc_avg / float(counter), acc_regime
def train(self, epoch):
self.model.train()
counter = 0
loss_avg = 0.0
loss_meta_avg = 0.0
acc_avg = 0.0
acc_multihead_avg = [0.0] * 3
for batch_data in tqdm(self.trainloader, f'Train epoch {epoch}'):
counter += 1
image, target, meta_target, structure_encoded, data_file = batch_data
image = renormalize(image)
image = image.cuda()
target = target.cuda()
if self.use_meta:
meta_target = meta_target.cuda()
model_outputs = self.model(image)
if len(model_outputs) == 3:
model_output, meta_pred, model_output_heads = model_outputs
else:
model_output, meta_pred = model_outputs
model_output_heads = None
loss = self.criterion(model_output, target)
loss_avg += loss.item()
acc = criteria.calculate_acc(model_output, target)
acc_avg += acc.item()
if self.use_meta:
if len(meta_pred.shape) > 2:
meta_pred = meta_pred[range(meta_pred.shape[0]), target]
loss_meta = self.criterion_meta(meta_pred, meta_target)
loss_meta_avg += loss_meta.item()
loss += self.use_meta * loss_meta
if self.args.multihead:
loss_head = [
self.criterion(output, target, reduction='none')
for output in model_output_heads
]
target_one_hot = torch.zeros_like(model_output_heads[0])
target_one_hot.scatter_(1, target.view(-1, 1), 1.0)
if self.args.multihead_mode is None:
weights = [1 / len(model_output_heads)
] * len(model_output_heads)
else:
probs = [
target_one_hot * output.detach().sigmoid() +
(1 - target_one_hot) * (1 - output.detach().sigmoid())
for output in model_output_heads
]
if self.args.multihead_mode == 'prob':
probs_sum = sum(probs)
weights = [prob / probs_sum for prob in probs]
elif self.args.multihead_mode == 'eprob':
e_probs = [prob.exp() for prob in probs]
e_probs_sum = sum(e_probs)
weights = [prob / e_probs_sum for prob in e_probs]
else:
raise ValueError(
f'Unsupported argument for multihead_mode: {self.args.multihead_mode}'
)
loss_multihead = sum([
weights[i] * loss_head[i] for i in range(len(loss_head))
]).mean()
loss += self.args.multihead_w * loss_multihead
acc_multihead = [
criteria.calculate_acc(x, target)
for x in model_output_heads
]
for i, x in enumerate(acc_multihead):
acc_multihead_avg[i] += x.item()
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
if self.use_meta:
print(
"Epoch {}, Train Avg Loss: {:.6f} META: {:.6f}, Train Avg Acc: {:.4f}"
.format(epoch, loss_avg / float(counter),
loss_meta_avg / float(counter),
acc_avg / float(counter)))
else:
print("Epoch {}, Train Avg Loss: {:.6f}, Train Avg Acc: {:.4f}".
format(epoch, loss_avg / float(counter),
acc_avg / float(counter)))
if self.args.multihead:
print(
f'Multihead: {[x / float(counter) for x in acc_multihead_avg]}'
)
return loss_avg / float(counter), acc_avg / float(counter)