def evaluate(self,
epoch: int = -1,
writer=None,
tag: str = 'validation') -> Tuple[str, bool]:
self.put_model_on_device()
total_error = []
# for batch in tqdm(self.data_loader.get_data_batch(), ascii=True, desc='evaluate'):
for batch in self.data_loader.get_data_batch():
with torch.no_grad():
predictions = self._net.forward(batch.observations,
train=False)
targets = data_to_tensor(batch.actions).type(
self._net.dtype).to(self._device)
error = self._criterion(predictions, targets).mean()
total_error.append(error)
error_distribution = Distribution(total_error)
self.put_model_back_to_original_device()
if writer is not None:
writer.write_distribution(error_distribution, tag)
if self._config.store_output_on_tensorboard and (epoch % 30 == 0
or tag == 'test'):
writer.write_output_image(predictions, f'{tag}/predictions')
writer.write_output_image(targets, f'{tag}/targets')
writer.write_output_image(torch.stack(batch.observations),
f'{tag}/inputs')
msg = f' {tag} {self._config.criterion} {error_distribution.mean: 0.3e} [{error_distribution.std:0.2e}]'
best_checkpoint = False
if self._lowest_validation_loss is None or error_distribution.mean < self._lowest_validation_loss:
self._lowest_validation_loss = error_distribution.mean
best_checkpoint = True
return msg, best_checkpoint
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
total_error = []
for batch in self.data_loader.sample_shuffled_batch():
self._optimizer.zero_grad()
targets = data_to_tensor(batch.actions).type(self._net.dtype).to(
self._device)
probabilities = self._net.forward_with_all_outputs(
batch.observations, train=True)
loss = self._criterion(probabilities[-1], targets).mean()
for index, prob in enumerate(probabilities[:-1]):
loss += self._criterion(prob, targets).mean()
loss.mean().backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
self._net.global_step += 1
total_error.append(loss.cpu().detach())
self.put_model_back_to_original_device()
error_distribution = Distribution(total_error)
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(error_distribution, 'training')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
for index, prob in enumerate(probabilities):
writer.write_output_image(prob,
f'training/predictions_{index}')
writer.write_output_image(targets, 'training/targets')
writer.write_output_image(torch.stack(batch.observations),
'training/inputs')
return f' training {self._config.criterion} {error_distribution.mean: 0.3e} [{error_distribution.std:0.2e}]'
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
total_error = []
task_error = []
domain_error = []
for source_batch, target_batch in zip(self.data_loader.sample_shuffled_batch(),
self.target_data_loader.sample_shuffled_batch()):
self._optimizer.zero_grad()
targets = data_to_tensor(source_batch.actions).type(self._net.dtype).to(self._device)
# task loss
predictions = self._net.forward(source_batch.observations, train=True)
task_loss = (1 - self._config.epsilon) * self._criterion(predictions, targets).mean()
# add domain adaptation loss
domain_loss = self._config.epsilon * self._domain_adaptation_criterion(
self._net.get_features(source_batch.observations, train=True),
self._net.get_features(target_batch.observations, train=True))
loss = task_loss + domain_loss
loss.backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
self._net.global_step += 1
task_error.append(task_loss.cpu().detach())
domain_error.append(domain_loss.cpu().detach())
total_error.append(loss.cpu().detach())
self.put_model_back_to_original_device()
if self._scheduler is not None:
self._scheduler.step()
task_error_distribution = Distribution(task_error)
domain_error_distribution = Distribution(domain_error)
total_error_distribution = Distribution(total_error)
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(task_error_distribution, 'training/task_error')
writer.write_distribution(domain_error_distribution, 'training/domain_error')
writer.write_distribution(total_error_distribution, 'training/total_error')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
writer.write_output_image(predictions, 'source/predictions')
writer.write_output_image(targets, 'source/targets')
writer.write_output_image(torch.stack(source_batch.observations), 'source/inputs')
writer.write_output_image(self._net.forward(target_batch.observations, train=True),
'target/predictions')
writer.write_output_image(torch.stack(target_batch.observations), 'target/inputs')
return f' training task: {self._config.criterion} {task_error_distribution.mean: 0.3e} ' \
f'[{task_error_distribution.std:0.2e}]' \
f' domain: {self._config.domain_adaptation_criterion} {domain_error_distribution.mean: 0.3e} ' \
f'[{domain_error_distribution.std:0.2e}]'
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
total_error = []
for batch in self.data_loader.sample_shuffled_batch():
self._optimizer.zero_grad()
targets = data_to_tensor(batch.actions).type(self._net.dtype).to(
self._device)
probabilities = self._net.forward_with_all_outputs(
batch.observations, train=True)
loss = self._criterion(probabilities[-1], targets).mean()
for index, prob in enumerate(probabilities[:-1]):
loss += self._criterion(prob, targets).mean()
loss.mean().backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
self._net.global_step += 1
total_error.append(loss.cpu().detach())
self.put_model_back_to_original_device()
error_distribution = Distribution(total_error)
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(error_distribution, 'training')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
for index, prob in enumerate(probabilities):
writer.write_output_image(prob,
f'training/predictions_{index}')
writer.write_output_image(targets, 'training/targets')
writer.write_output_image(torch.stack(batch.observations),
'training/inputs')
if self._config.store_feature_maps_on_tensorboard and epoch % 30 == 0:
outputs = self._net.forward_with_intermediate_outputs(
batch.observations, train=False)
for i in range(4): # store first 5 images of batch
for layer in ['x1', 'x2', 'x3', 'x4']:
feature_maps = outputs[layer][i].flatten(start_dim=0,
end_dim=0)
title = f'feature_map/layer_{layer}/{i}'
# title += 'inds_' + '_'.join([str(v.item()) for v in winning_indices.indices])
# title += '_vals_' + '_'.join([f'{v.item():0.2f}' for v in winning_indices.values])
writer.write_output_image(feature_maps, title)
writer.write_figure(tag='gradient',
figure=plot_gradient_flow(
self._net.named_parameters()))
return f' training {self._config.criterion} {error_distribution.mean: 0.3e} [{error_distribution.std:0.2e}]'
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
total_error = []
# for batch in tqdm(self.data_loader.sample_shuffled_batch(), ascii=True, desc='train'):
for batch in self.data_loader.sample_shuffled_batch():
self._optimizer.zero_grad()
targets = data_to_tensor(batch.actions).type(self._net.dtype).to(
self._device)
if self._config.add_KL_divergence_loss:
predictions, mean, std = self._net.forward_with_distribution(
batch.observations, train=True)
else:
predictions = self._net.forward(batch.observations, train=True)
loss = self._criterion(predictions, targets).mean()
if self._config.add_KL_divergence_loss:
# https://arxiv.org/pdf/1312.6114.pdf
KL_loss = -0.5 * torch.sum(1 + std.pow(2).log() - mean.pow(2) -
std.pow(2))
loss += KL_loss
loss.backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
self._net.global_step += 1
total_error.append(loss.cpu().detach())
self.put_model_back_to_original_device()
if self._scheduler is not None:
self._scheduler.step()
error_distribution = Distribution(total_error)
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(error_distribution, 'training')
if self._config.add_KL_divergence_loss:
writer.write_scalar(KL_loss, 'KL_divergence')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
writer.write_output_image(predictions, 'training/predictions')
writer.write_output_image(targets, 'training/targets')
writer.write_output_image(torch.stack(batch.observations),
'training/inputs')
return f' training {self._config.criterion} {error_distribution.mean: 0.3e} [{error_distribution.std:0.2e}]'
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
total_error = []
task_error = []
domain_error = []
for source_batch, target_batch in zip(self.data_loader.sample_shuffled_batch(),
self.target_data_loader.sample_shuffled_batch()):
self._optimizer.zero_grad()
targets = data_to_tensor(source_batch.actions).type(self._net.dtype).to(self._device)
# deep supervision loss
probabilities = self._net.forward_with_all_outputs(source_batch.observations, train=True)
task_loss = self._criterion(probabilities[-1], targets).mean()
for index, prob in enumerate(probabilities[:-1]):
task_loss += self._criterion(prob, targets).mean()
task_loss *= (1 - self._config.epsilon)
# add domain adaptation loss on distribution of output pixels at each output
domain_loss = sum([self._domain_adaptation_criterion(sp.flatten().unsqueeze(1), tp.flatten().unsqueeze(1))
for sp, tp in zip(self._net.forward_with_all_outputs(source_batch.observations,
train=True),
self._net.forward_with_all_outputs(target_batch.observations,
train=True))
]) * self._config.epsilon
loss = task_loss + domain_loss
loss.backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
self._net.global_step += 1
task_error.append(task_loss.cpu().detach().numpy())
domain_error.append(domain_loss.cpu().detach().numpy())
total_error.append(loss.cpu().detach().numpy())
self.put_model_back_to_original_device()
if self._scheduler is not None:
self._scheduler.step()
task_error_distribution = Distribution(task_error)
domain_error_distribution = Distribution(domain_error)
total_error_distribution = Distribution(total_error)
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(task_error_distribution, 'training/task_error')
writer.write_distribution(domain_error_distribution, 'training/domain_error')
writer.write_distribution(total_error_distribution, 'training/total_error')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
writer.write_output_image(probabilities[-1], 'source/predictions')
writer.write_output_image(targets, 'source/targets')
writer.write_output_image(torch.stack(source_batch.observations), 'source/inputs')
writer.write_output_image(self._net.forward(target_batch.observations, train=False),
'target/predictions')
writer.write_output_image(torch.stack(target_batch.observations), 'target/inputs')
if self._config.store_feature_maps_on_tensorboard and epoch % 30 == 0:
for name, batch in zip(['source', 'target'], [source_batch, target_batch]):
outputs = self._net.forward_with_intermediate_outputs(batch.observations, train=False)
for i in range(4): # store first 5 images of batch
for layer in ['x1', 'x2', 'x3', 'x4']:
feature_maps = outputs[layer][i].flatten(start_dim=0, end_dim=0)
title = f'feature_map/{name}/layer_{layer}/{i}'
# title += 'inds_' + '_'.join([str(v.item()) for v in winning_indices.indices])
# title += '_vals_' + '_'.join([f'{v.item():0.2f}' for v in winning_indices.values])
writer.write_output_image(feature_maps, title)
return f' task {self._config.criterion} ' \
f'{task_error_distribution.mean: 0.3e} ' \
f'[{task_error_distribution.std:0.2e}] ' \
f' domain {self._config.domain_adaptation_criterion} ' \
f'{domain_error_distribution.mean: 0.3e} ' \
f'[{domain_error_distribution.std: 0.2e}]'
def _train_main_network(self, epoch: int = -1, writer=None) -> str:
deeply_supervised_error = []
discriminator_error = []
for sim_batch, real_batch in zip(
self.data_loader.sample_shuffled_batch(),
self.target_data_loader.sample_shuffled_batch()):
self._optimizer.zero_grad()
# normal deep supervision loss
targets = data_to_tensor(sim_batch.actions).type(
self._net.dtype).to(self._device)
probabilities = self._net.forward_with_all_outputs(
sim_batch.observations, train=True)
loss = self._criterion(probabilities[-1], targets).mean()
for index, prob in enumerate(probabilities[:-1]):
loss += self._criterion(prob, targets).mean()
deeply_supervised_error.append(loss.mean().cpu().detach())
# adversarial loss on discriminator data
network_outputs = torch.cat(
self._net.forward_with_all_outputs(
real_batch.observations, train=True)).unsqueeze(dim=1)
discriminator_loss = self._net.discriminate(network_outputs,
train=False).mean()
#results = self._net.forward_with_intermediate_outputs(real_batch.observations, train=True)
#feature_maps =
# combine losses with epsilon weight
loss *= (1 - self._config.epsilon)
loss += self._config.epsilon * discriminator_loss
loss.mean().backward()
discriminator_error.append(
discriminator_loss.mean().cpu().detach())
# clip gradients
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
supervised_error_distribution = Distribution(deeply_supervised_error)
discriminator_error_distribution = Distribution(discriminator_error)
if writer is not None:
writer.write_distribution(supervised_error_distribution,
'training_loss_from_deep_supervision')
writer.write_distribution(discriminator_error_distribution,
'training_loss_from_discriminator')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
for index, prob in enumerate(probabilities):
writer.write_output_image(prob,
f'training/predictions_{index}')
writer.write_output_image(targets, 'training/targets')
writer.write_output_image(torch.stack(sim_batch.observations),
'training/inputs')
for index, prob in enumerate(
self._net.forward_with_all_outputs(real_batch.observations,
train=False)):
writer.write_output_image(prob, f'real/predictions_{index}')
writer.write_output_image(torch.stack(real_batch.observations),
'real/inputs')
return f' Training: supervision {self._config.criterion} {supervised_error_distribution.mean: 0.3e} ' \
f'[{supervised_error_distribution.std:0.2e}]' \
f' discriminator {discriminator_error_distribution.mean: 0.3e} ' \
f'[{discriminator_error_distribution.std:0.2e}]'