def test_big_data_hdf5_loop(self):
# create 3 datasets as hdf5 files
hdf5_files = []
infos = []
for index in range(3):
output_path = os.path.join(self.output_dir, f'ds{index}')
os.makedirs(output_path, exist_ok=True)
config_dict = {
'output_path': output_path,
'store_hdf5': True,
'training_validation_split': 1.0
}
config = DataSaverConfig().create(config_dict=config_dict)
self.data_saver = DataSaver(config=config)
infos.append(
generate_dummy_dataset(self.data_saver,
num_runs=2,
input_size=(3, 10, 10),
fixed_input_value=(0.3 * index) *
np.ones((3, 10, 10)),
store_hdf5=True))
self.assertTrue(
os.path.isfile(os.path.join(output_path, 'train.hdf5')))
hdf5_files.append(os.path.join(output_path, 'train.hdf5'))
hdf5_files.append(os.path.join(output_path, 'wrong.hdf5'))
# create data loader with big data tag and three hdf5 training sets
conf = {
'output_path': self.output_dir,
'hdf5_files': hdf5_files,
'batch_size': 15,
'loop_over_hdf5_files': True
}
loader = DataLoader(DataLoaderConfig().create(config_dict=conf))
# sample data batches and see that index increases every two batches sampled
for batch in loader.get_data_batch():
self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0)
for batch in loader.get_data_batch():
self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0.3,
2)
for batch in loader.get_data_batch():
self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0.6,
2)
for batch in loader.get_data_batch():
self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0, 2)
for batch in loader.sample_shuffled_batch():
self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0.3,
2)
for batch in loader.sample_shuffled_batch():
self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0.6,
2)
for batch in loader.sample_shuffled_batch():
self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0, 2)
def test_sample_batch(self):
self.info = generate_dummy_dataset(self.data_saver,
num_runs=20,
input_size=(100, 100, 3),
output_size=(3, ),
continuous=False)
max_num_batches = 2
config_dict = {
'data_directories': self.info['episode_directories'],
'output_path': self.output_dir,
'random_seed': 1,
'batch_size': 3
}
data_loader = DataLoader(config=DataLoaderConfig().create(
config_dict=config_dict))
data_loader.load_dataset()
first_batch = []
index = 0
for index, batch in enumerate(
data_loader.sample_shuffled_batch(
max_number_of_batches=max_num_batches)):
if index == 0:
first_batch = deepcopy(batch)
self.assertEqual(len(batch), config_dict['batch_size'])
self.assertEqual(index, max_num_batches - 1)
# test sampling seed for reproduction
config_dict['random_seed'] = 2
data_loader = DataLoader(config=DataLoaderConfig().create(
config_dict=config_dict))
data_loader.load_dataset()
second_batch = []
for index, batch in enumerate(
data_loader.sample_shuffled_batch(
max_number_of_batches=max_num_batches)):
second_batch = deepcopy(batch)
break
self.assertNotEqual(np.sum(np.asarray(first_batch.observations[0])),
np.sum(np.asarray(second_batch.observations[0])))
config_dict['random_seed'] = 1
data_loader = DataLoader(config=DataLoaderConfig().create(
config_dict=config_dict))
data_loader.load_dataset()
third_batch = []
for index, batch in enumerate(
data_loader.sample_shuffled_batch(
max_number_of_batches=max_num_batches)):
third_batch = deepcopy(batch)
break
self.assertEqual(np.sum(np.asarray(first_batch.observations[0])),
np.sum(np.asarray(third_batch.observations[0])))
class DomainAdaptationTrainer(Trainer):
def __init__(self, config: TrainerConfig, network: BaseNet, quiet: bool = False):
super().__init__(config, network, quiet=True)
self._config.epsilon = 0.2 if self._config.epsilon == "default" else self._config.epsilon
self.target_data_loader = DataLoader(config=self._config.target_data_loader_config)
self.target_data_loader.load_dataset()
self._domain_adaptation_criterion = eval(f'{self._config.domain_adaptation_criterion}()') \
if not self._config.domain_adaptation_criterion == 'default' else MMDLossZhao()
self._domain_adaptation_criterion.to(self._device)
if not quiet:
self._optimizer = eval(f'torch.optim.{self._config.optimizer}')(params=self._net.parameters(),
lr=self._config.learning_rate,
weight_decay=self._config.weight_decay)
lambda_function = lambda f: 1 - f / self._config.scheduler_config.number_of_epochs
self._scheduler = torch.optim.lr_scheduler.LambdaLR(self._optimizer, lr_lambda=lambda_function) \
if self._config.scheduler_config is not None else None
self._logger = get_logger(name=get_filename_without_extension(__file__),
output_path=config.output_path,
quiet=False)
cprint(f'Started.', self._logger)
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}]'