def test_poutyne(): callback = PlotLossesPoutyne(outputs=(CheckOutput(), )) network = Network() optimizer = optim.Adam(params=network.parameters(), lr=0.001) loss_fn = nn.CrossEntropyLoss() train_dataloader = get_random_data() model = Model(network, optimizer, loss_fn) model.fit_generator(train_dataloader, epochs=2, callbacks=[callback])
class LRSchedulersTest(TestCase): batch_size = 20 epochs = 10 steps_per_epoch = 5 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) self.train_gen = some_data_generator(20) self.valid_gen = some_data_generator(20) def test_lambda_lr_integration(self): my_lambda = lambda epoch: 0.95**epoch lambda_lr = LambdaLR(lr_lambda=[my_lambda]) self._fit_with_callback_integration(lambda_lr) def test_step_lr_integration(self): step_lr = StepLR(step_size=3) self._fit_with_callback_integration(step_lr) def test_multistep_lr_integration(self): multistep_lr = MultiStepLR(milestones=[2, 5, 7]) self._fit_with_callback_integration(multistep_lr) def test_exponential_lr_integration(self): exponential_lr = ExponentialLR(gamma=0.01) self._fit_with_callback_integration(exponential_lr) def test_cosine_annealing_lr_integration(self): cosine_annealing_lr = CosineAnnealingLR(T_max=8) self._fit_with_callback_integration(cosine_annealing_lr) def test_reduce_lr_on_plateau_integration(self): reduce_lr = ReduceLROnPlateau(monitor='loss', patience=3) self._fit_with_callback_integration(reduce_lr) def _fit_with_callback_integration(self, callback): self.model.fit_generator( self.train_gen, self.valid_gen, epochs=LRSchedulersTest.epochs, steps_per_epoch=LRSchedulersTest.steps_per_epoch, callbacks=[callback], ) def test_exception_is_thrown_on_optimizer_argument(self): with self.assertRaises(ValueError): StepLR(self.optimizer, step_size=3)
def test_correct_optim_calls_1_batch_per_step(self): train_generator = some_data_tensor_generator(ModelTest.batch_size) mocked_optimizer = some_mocked_optimizer() mocked_optim_model = Model(self.pytorch_network, mocked_optimizer, self.loss_function, batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics) mocked_optim_model.fit_generator(train_generator, None, epochs=1, steps_per_epoch=1, batches_per_step=1) self.assertEqual(1, mocked_optimizer.step.call_count) self.assertEqual(1, mocked_optimizer.zero_grad.call_count)
def test_metrics_integration(self): num_steps = 10 model = Model(self.pytorch_network, self.optimizer, self.loss_function, batch_metrics=[F.mse_loss]) train_generator = some_data_tensor_generator(ModelTest.batch_size) valid_generator = some_data_tensor_generator(ModelTest.batch_size) model.fit_generator(train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch, callbacks=[self.mock_callback]) generator = some_data_tensor_generator(ModelTest.batch_size) loss, mse = model.evaluate_generator(generator, steps=num_steps) self.assertEqual(type(loss), float) self.assertEqual(type(mse), float)
def test_correct_optim_calls__valid_n_batches_per_step(self): n_batches = 5 items_per_batch = int(ModelTest.batch_size / n_batches) x = torch.rand(n_batches, items_per_batch, 1) y = torch.rand(n_batches, items_per_batch, 1) mocked_optimizer = some_mocked_optimizer() mocked_optim_model = Model(self.pytorch_network, mocked_optimizer, self.loss_function, batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics) mocked_optim_model.fit_generator(list(zip(x, y)), None, epochs=1, batches_per_step=n_batches) self.assertEqual(1, mocked_optimizer.step.call_count) self.assertEqual(1, mocked_optimizer.zero_grad.call_count)
class BaseTensorBoardLoggerTest: SummaryWriter = None batch_size = 20 lr = 1e-3 num_epochs = 10 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=BaseTensorBoardLoggerTest.lr) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) self.temp_dir_obj = TemporaryDirectory() # pylint: disable=not-callable self.writer = self.SummaryWriter(self.temp_dir_obj.name) self.writer.add_scalars = MagicMock() def tearDown(self): self.temp_dir_obj.cleanup() def test_logging(self): train_gen = some_data_generator(20) valid_gen = some_data_generator(20) logger = TensorBoardLogger(self.writer) history = self.model.fit_generator( train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[logger] ) self._test_logging(history) def test_multiple_learning_rates(self): train_gen = some_data_generator(20) valid_gen = some_data_generator(20) logger = TensorBoardLogger(self.writer) lrs = [BaseCSVLoggerTest.lr, BaseCSVLoggerTest.lr / 2] optimizer = torch.optim.SGD( [dict(params=[self.pytorch_network.weight], lr=lrs[0]), dict(params=[self.pytorch_network.bias], lr=lrs[1])] ) model = Model(self.pytorch_network, optimizer, self.loss_function) history = model.fit_generator( train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[logger] ) self._test_logging(history, lrs=lrs) def _test_logging(self, history, lrs=None): if lrs is None: lrs = [BaseCSVLoggerTest.lr] calls = [] for h in history: calls.append(call('loss', {'loss': h['loss'], 'val_loss': h['val_loss']}, h['epoch'])) if len(lrs) == 1: calls.append(call('lr', {'lr': self.lr}, h['epoch'])) else: calls.append(call('lr', {f'lr_group_{i}': lr for i, lr in enumerate(lrs)}, h['epoch'])) self.writer.add_scalars.assert_has_calls(calls, any_order=True)
def test_multiple_learning_rates(self): train_gen = some_data_generator(20) valid_gen = some_data_generator(20) logger = self.CSVLogger(self.csv_filename) lrs = [BaseCSVLoggerTest.lr, BaseCSVLoggerTest.lr / 2] optimizer = torch.optim.SGD( [dict(params=[self.pytorch_network.weight], lr=lrs[0]), dict(params=[self.pytorch_network.bias], lr=lrs[1])] ) model = Model(self.pytorch_network, optimizer, self.loss_function) history = model.fit_generator( train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[logger] ) self._test_logging(history, lrs=lrs)
def test_epoch_metrics_integration(self): model = Model(self.pytorch_network, self.optimizer, self.loss_function, epoch_metrics=[SomeEpochMetric()]) train_generator = some_data_tensor_generator(ModelTest.batch_size) valid_generator = some_data_tensor_generator(ModelTest.batch_size) logs = model.fit_generator(train_generator, valid_generator, epochs=1, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch) actual_value = logs[-1]['some_epoch_metric'] val_actual_value = logs[-1]['val_some_epoch_metric'] expected_value = 5 self.assertEqual(val_actual_value, expected_value) self.assertEqual(actual_value, expected_value)
class BaseTensorBoardLoggerTest: SummaryWriter = None batch_size = 20 lr = 1e-3 num_epochs = 10 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=BaseTensorBoardLoggerTest.lr) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) self.temp_dir_obj = TemporaryDirectory() # pylint: disable=not-callable self.writer = self.SummaryWriter(self.temp_dir_obj.name) self.writer.add_scalars = MagicMock() def tearDown(self): self.temp_dir_obj.cleanup() def test_logging(self): train_gen = some_data_generator(20) valid_gen = some_data_generator(20) logger = TensorBoardLogger(self.writer) history = self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[logger]) self._test_logging(history) def _test_logging(self, history): calls = list() for h in history: calls.append( call('loss', { 'loss': h['loss'], 'val_loss': h['val_loss'] }, h['epoch'])) calls.append(call('lr', {'lr': self.lr}, h['epoch'])) self.writer.add_scalars.assert_has_calls(calls, any_order=True)
class EarlyStoppingTest(TestCase): batch_size = 20 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) def test_integration(self): train_gen = some_data_generator(20) valid_gen = some_data_generator(20) earlystopper = EarlyStopping(monitor='val_loss', min_delta=0, patience=2, verbose=False) self.model.fit_generator(train_gen, valid_gen, epochs=10, steps_per_epoch=5, callbacks=[earlystopper]) def test_early_stopping_patience_of_1(self): earlystopper = EarlyStopping(monitor='val_loss', min_delta=0, patience=1, verbose=False) val_losses = [8, 4, 5, 2] early_stop_epoch = 3 self._test_early_stopping(earlystopper, val_losses, early_stop_epoch) def test_early_stopping_with_delta(self): earlystopper = EarlyStopping(monitor='val_loss', min_delta=3, patience=2, verbose=False) val_losses = [8, 4, 5, 2, 2] early_stop_epoch = 4 self._test_early_stopping(earlystopper, val_losses, early_stop_epoch) def test_early_stopping_with_max(self): earlystopper = EarlyStopping(monitor='val_loss', mode='max', min_delta=0, patience=2, verbose=False) val_losses = [2, 8, 4, 5, 2] early_stop_epoch = 4 self._test_early_stopping(earlystopper, val_losses, early_stop_epoch) def _test_early_stopping(self, earlystopper, val_losses, early_stop_epoch): generator = some_data_generator(EarlyStoppingTest.batch_size) self.model.stop_training = False earlystopper.set_params({'epochs': len(val_losses), 'steps': 1}) earlystopper.set_model(self.model) earlystopper.on_train_begin({}) for epoch, val_loss in enumerate(val_losses, 1): earlystopper.on_epoch_begin(epoch, {}) earlystopper.on_train_batch_begin(1, {}) loss = self._update_model(generator) earlystopper.on_train_batch_end(1, {'batch': 1, 'size': EarlyStoppingTest.batch_size, 'loss': loss}) earlystopper.on_epoch_end(epoch, {'epoch': epoch, 'loss': loss, 'val_loss': val_loss}) self.assertEqual(self.model.stop_training, epoch == early_stop_epoch) if epoch == early_stop_epoch: break earlystopper.on_train_end({}) def _update_model(self, generator): self.pytorch_network.zero_grad() x, y = next(generator) pred_y = self.pytorch_network(x) loss = self.loss_function(pred_y, y) loss.backward() self.optimizer.step() return float(loss)
class ModelTest(ModelFittingTestCase): # pylint: disable=too-many-public-methods def setUp(self): super().setUp() torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.Adam(self.pytorch_network.parameters(), lr=1e-3) self.batch_metrics = [ some_batch_metric_1, ('custom_name', some_batch_metric_2), repeat_batch_metric, repeat_batch_metric ] self.batch_metrics_names = [ 'some_batch_metric_1', 'custom_name', 'repeat_batch_metric1', 'repeat_batch_metric2' ] self.batch_metrics_values = [ some_metric_1_value, some_metric_2_value, repeat_batch_metric_value, repeat_batch_metric_value ] self.epoch_metrics = [SomeConstantEpochMetric()] self.epoch_metrics_names = ['some_constant_epoch_metric'] self.epoch_metrics_values = [some_constant_epoch_metric_value] self.model = Model(self.pytorch_network, self.optimizer, self.loss_function, batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics) def test_fitting_tensor_generator(self): train_generator = some_data_tensor_generator(ModelTest.batch_size) valid_generator = some_data_tensor_generator(ModelTest.batch_size) logs = self.model.fit_generator( train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch, callbacks=[self.mock_callback]) params = { 'epochs': ModelTest.epochs, 'steps': ModelTest.steps_per_epoch, 'valid_steps': ModelTest.steps_per_epoch } self._test_callbacks_train(params, logs, valid_steps=ModelTest.steps_per_epoch) def test_fitting_without_valid_generator(self): train_generator = some_data_tensor_generator(ModelTest.batch_size) logs = self.model.fit_generator( train_generator, None, epochs=ModelTest.epochs, steps_per_epoch=ModelTest.steps_per_epoch, callbacks=[self.mock_callback]) params = { 'epochs': ModelTest.epochs, 'steps': ModelTest.steps_per_epoch } self._test_callbacks_train(params, logs, has_valid=False) def test_correct_optim_calls_1_batch_per_step(self): train_generator = some_data_tensor_generator(ModelTest.batch_size) mocked_optimizer = some_mocked_optimizer() mocked_optim_model = Model(self.pytorch_network, mocked_optimizer, self.loss_function, batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics) mocked_optim_model.fit_generator(train_generator, None, epochs=1, steps_per_epoch=1, batches_per_step=1) self.assertEqual(1, mocked_optimizer.step.call_count) self.assertEqual(1, mocked_optimizer.zero_grad.call_count) def test_correct_optim_calls__valid_n_batches_per_step(self): n_batches = 5 items_per_batch = int(ModelTest.batch_size / n_batches) x = torch.rand(n_batches, items_per_batch, 1) y = torch.rand(n_batches, items_per_batch, 1) mocked_optimizer = some_mocked_optimizer() mocked_optim_model = Model(self.pytorch_network, mocked_optimizer, self.loss_function, batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics) mocked_optim_model.fit_generator(list(zip(x, y)), None, epochs=1, batches_per_step=n_batches) self.assertEqual(1, mocked_optimizer.step.call_count) self.assertEqual(1, mocked_optimizer.zero_grad.call_count) def test_fitting_generator_n_batches_per_step(self): total_batch_size = 6 x = torch.rand(1, total_batch_size, 1) y = torch.rand(1, total_batch_size, 1) initial_params = self.model.get_weight_copies() self.model.fit_generator(list(zip(x, y)), None, epochs=1, batches_per_step=1) expected_params = self.model.get_weight_copies() for mini_batch_size in [1, 2, 5]: self.model.set_weights(initial_params) n_batches_per_step = int(total_batch_size / mini_batch_size) x.resize_((n_batches_per_step, mini_batch_size, 1)) y.resize_((n_batches_per_step, mini_batch_size, 1)) self.model.fit_generator(list(zip(x, y)), None, epochs=1, batches_per_step=n_batches_per_step) returned_params = self.model.get_weight_copies() self.assertEqual(returned_params.keys(), expected_params.keys()) for k in expected_params.keys(): np.testing.assert_almost_equal(returned_params[k].numpy(), expected_params[k].numpy(), decimal=4) def test_fitting_generator_n_batches_per_step_higher_than_num_batches( self): total_batch_size = 6 x = torch.rand(1, total_batch_size, 1) y = torch.rand(1, total_batch_size, 1) initial_params = self.model.get_weight_copies() self.model.fit_generator(list(zip(x, y)), None, epochs=1, batches_per_step=1) expected_params = self.model.get_weight_copies() self.model.set_weights(initial_params) self.model.fit_generator(list(zip(x, y)), None, epochs=1, batches_per_step=2) returned_params = self.model.get_weight_copies() self.assertEqual(returned_params.keys(), expected_params.keys()) for k in expected_params.keys(): np.testing.assert_almost_equal(returned_params[k].numpy(), expected_params[k].numpy(), decimal=4) def test_fitting_generator_n_batches_per_step_uneven_batches(self): total_batch_size = 6 x = torch.rand(1, total_batch_size, 1) y = torch.rand(1, total_batch_size, 1) initial_params = self.model.get_weight_copies() self.model.fit_generator(list(zip(x, y)), None, epochs=1, batches_per_step=1) expected_params = self.model.get_weight_copies() x.squeeze_(dim=0) y.squeeze_(dim=0) uneven_chunk_sizes = [4, 5] for chunk_size in uneven_chunk_sizes: self.model.set_weights(initial_params) splitted_x = x.split(chunk_size) splitted_y = y.split(chunk_size) n_batches_per_step = ceil(total_batch_size / chunk_size) self.model.fit_generator(list(zip(splitted_x, splitted_y)), None, epochs=1, batches_per_step=n_batches_per_step) returned_params = self.model.get_weight_copies() self.assertEqual(returned_params.keys(), expected_params.keys()) for k in expected_params.keys(): np.testing.assert_almost_equal(returned_params[k].numpy(), expected_params[k].numpy(), decimal=4) def test_fitting_ndarray_generator(self): train_generator = some_ndarray_generator(ModelTest.batch_size) valid_generator = some_ndarray_generator(ModelTest.batch_size) logs = self.model.fit_generator( train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch, callbacks=[self.mock_callback]) params = { 'epochs': ModelTest.epochs, 'steps': ModelTest.steps_per_epoch, 'valid_steps': ModelTest.steps_per_epoch } self._test_callbacks_train(params, logs, valid_steps=ModelTest.steps_per_epoch) def test_fitting_with_data_loader(self): train_real_steps_per_epoch = 30 train_batch_size = ModelTest.batch_size train_final_batch_missing_samples = 7 train_size = train_real_steps_per_epoch * train_batch_size - \ train_final_batch_missing_samples train_x = torch.rand(train_size, 1) train_y = torch.rand(train_size, 1) train_dataset = TensorDataset(train_x, train_y) train_generator = DataLoader(train_dataset, train_batch_size) valid_real_steps_per_epoch = 10 valid_batch_size = 15 valid_final_batch_missing_samples = 3 valid_size = valid_real_steps_per_epoch * valid_batch_size - \ valid_final_batch_missing_samples valid_x = torch.rand(valid_size, 1) valid_y = torch.rand(valid_size, 1) valid_dataset = TensorDataset(valid_x, valid_y) valid_generator = DataLoader(valid_dataset, valid_batch_size) logs = self.model.fit_generator(train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=None, validation_steps=None, callbacks=[self.mock_callback]) params = { 'epochs': ModelTest.epochs, 'steps': train_real_steps_per_epoch, 'valid_steps': valid_real_steps_per_epoch } self._test_callbacks_train(params, logs) def test_fitting_generator_calls(self): train_real_steps_per_epoch = 30 train_batch_size = ModelTest.batch_size train_final_batch_missing_samples = 7 train_size = train_real_steps_per_epoch * train_batch_size - \ train_final_batch_missing_samples train_x = torch.rand(train_size, 1) train_y = torch.rand(train_size, 1) train_dataset = TensorDataset(train_x, train_y) train_generator = DataLoader(train_dataset, train_batch_size) valid_real_steps_per_epoch = 10 valid_batch_size = 15 valid_final_batch_missing_samples = 3 valid_size = valid_real_steps_per_epoch * valid_batch_size - \ valid_final_batch_missing_samples valid_x = torch.rand(valid_size, 1) valid_y = torch.rand(valid_size, 1) valid_dataset = TensorDataset(valid_x, valid_y) valid_generator = DataLoader(valid_dataset, valid_batch_size) mock_train_generator = IterableMock(train_generator) mock_valid_generator = IterableMock(valid_generator) self.model.fit_generator(mock_train_generator, mock_valid_generator, epochs=ModelTest.epochs) expected_train_calls = ['__len__'] + \ (['__iter__'] + ['__next__'] * train_real_steps_per_epoch) * ModelTest.epochs expected_valid_calls = ['__len__'] + \ (['__iter__'] + ['__next__'] * valid_real_steps_per_epoch) * ModelTest.epochs self.assertEqual(mock_train_generator.calls, expected_train_calls) self.assertEqual(mock_valid_generator.calls, expected_valid_calls) def test_fitting_generator_calls_with_longer_validation_set(self): train_real_steps_per_epoch = 30 train_batch_size = ModelTest.batch_size train_final_batch_missing_samples = 7 train_size = train_real_steps_per_epoch * train_batch_size - \ train_final_batch_missing_samples train_x = torch.rand(train_size, 1) train_y = torch.rand(train_size, 1) train_dataset = TensorDataset(train_x, train_y) train_generator = DataLoader(train_dataset, train_batch_size) valid_real_steps_per_epoch = 40 valid_batch_size = 15 valid_final_batch_missing_samples = 3 valid_size = valid_real_steps_per_epoch * valid_batch_size - \ valid_final_batch_missing_samples valid_x = torch.rand(valid_size, 1) valid_y = torch.rand(valid_size, 1) valid_dataset = TensorDataset(valid_x, valid_y) valid_generator = DataLoader(valid_dataset, valid_batch_size) mock_train_generator = IterableMock(train_generator) mock_valid_generator = IterableMock(valid_generator) self.model.fit_generator(mock_train_generator, mock_valid_generator, epochs=ModelTest.epochs) expected_train_calls = ['__len__'] + \ (['__iter__'] + ['__next__'] * train_real_steps_per_epoch) * ModelTest.epochs expected_valid_calls = ['__len__'] + \ (['__iter__'] + ['__next__'] * valid_real_steps_per_epoch) * ModelTest.epochs self.assertEqual(mock_train_generator.calls, expected_train_calls) self.assertEqual(mock_valid_generator.calls, expected_valid_calls) def test_fitting_with_tensor(self): train_real_steps_per_epoch = 30 train_batch_size = ModelTest.batch_size train_final_batch_missing_samples = 7 train_size = train_real_steps_per_epoch * train_batch_size - \ train_final_batch_missing_samples train_x = torch.rand(train_size, 1) train_y = torch.rand(train_size, 1) valid_real_steps_per_epoch = 10 # valid_batch_size will be the same as train_batch_size in the fit method. valid_batch_size = train_batch_size valid_final_batch_missing_samples = 3 valid_size = valid_real_steps_per_epoch * valid_batch_size - \ valid_final_batch_missing_samples valid_x = torch.rand(valid_size, 1) valid_y = torch.rand(valid_size, 1) logs = self.model.fit(train_x, train_y, validation_data=(valid_x, valid_y), epochs=ModelTest.epochs, batch_size=train_batch_size, steps_per_epoch=None, validation_steps=None, callbacks=[self.mock_callback]) params = { 'epochs': ModelTest.epochs, 'steps': train_real_steps_per_epoch, 'valid_steps': valid_real_steps_per_epoch } self._test_callbacks_train(params, logs) def test_fitting_with_np_array(self): train_real_steps_per_epoch = 30 train_batch_size = ModelTest.batch_size train_final_batch_missing_samples = 7 train_size = train_real_steps_per_epoch * train_batch_size - \ train_final_batch_missing_samples train_x = np.random.rand(train_size, 1).astype(np.float32) train_y = np.random.rand(train_size, 1).astype(np.float32) valid_real_steps_per_epoch = 10 # valid_batch_size will be the same as train_batch_size in the fit method. valid_batch_size = train_batch_size valid_final_batch_missing_samples = 3 valid_size = valid_real_steps_per_epoch * valid_batch_size - \ valid_final_batch_missing_samples valid_x = np.random.rand(valid_size, 1).astype(np.float32) valid_y = np.random.rand(valid_size, 1).astype(np.float32) logs = self.model.fit(train_x, train_y, validation_data=(valid_x, valid_y), epochs=ModelTest.epochs, batch_size=train_batch_size, steps_per_epoch=None, validation_steps=None, callbacks=[self.mock_callback]) params = { 'epochs': ModelTest.epochs, 'steps': train_real_steps_per_epoch, 'valid_steps': valid_real_steps_per_epoch } self._test_callbacks_train(params, logs) def test_fitting_with_generator_with_len(self): train_real_steps_per_epoch = 30 train_generator = SomeDataGeneratorWithLen( batch_size=ModelTest.batch_size, length=train_real_steps_per_epoch, num_missing_samples=7) valid_real_steps_per_epoch = 10 valid_generator = SomeDataGeneratorWithLen( batch_size=15, length=valid_real_steps_per_epoch, num_missing_samples=3) logs = self.model.fit_generator(train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=None, validation_steps=None, callbacks=[self.mock_callback]) params = { 'epochs': ModelTest.epochs, 'steps': train_real_steps_per_epoch, 'valid_steps': valid_real_steps_per_epoch } self._test_callbacks_train(params, logs) def test_fitting_with_generator_with_stop_iteration(self): train_real_steps_per_epoch = 30 train_generator = SomeDataGeneratorUsingStopIteration( batch_size=ModelTest.batch_size, length=train_real_steps_per_epoch) valid_generator = SomeDataGeneratorUsingStopIteration(batch_size=15, length=10) logs = self.model.fit_generator(train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=None, validation_steps=None, callbacks=[self.mock_callback]) params = {'epochs': ModelTest.epochs, 'steps': None} self._test_callbacks_train(params, logs, steps=train_real_steps_per_epoch) def test_tensor_train_on_batch(self): x = torch.rand(ModelTest.batch_size, 1) y = torch.rand(ModelTest.batch_size, 1) loss, metrics = self.model.train_on_batch(x, y) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values) def test_train_on_batch_with_pred(self): x = torch.rand(ModelTest.batch_size, 1) y = torch.rand(ModelTest.batch_size, 1) loss, metrics, pred_y = self.model.train_on_batch(x, y, return_pred=True) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values) self.assertEqual(pred_y.shape, (ModelTest.batch_size, 1)) def test_ndarray_train_on_batch(self): x = np.random.rand(ModelTest.batch_size, 1).astype(np.float32) y = np.random.rand(ModelTest.batch_size, 1).astype(np.float32) loss, metrics = self.model.train_on_batch(x, y) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values) def test_evaluate(self): x = torch.rand(ModelTest.evaluate_dataset_len, 1) y = torch.rand(ModelTest.evaluate_dataset_len, 1) loss, metrics = self.model.evaluate(x, y, batch_size=ModelTest.batch_size) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values + self.epoch_metrics_values) def test_evaluate_with_pred(self): x = torch.rand(ModelTest.evaluate_dataset_len, 1) y = torch.rand(ModelTest.evaluate_dataset_len, 1) # We also test the unpacking. _, _, pred_y = self.model.evaluate(x, y, batch_size=ModelTest.batch_size, return_pred=True) self.assertEqual(pred_y.shape, (ModelTest.evaluate_dataset_len, 1)) def test_evaluate_with_callback(self): x = torch.rand(ModelTest.evaluate_dataset_len, 1) y = torch.rand(ModelTest.evaluate_dataset_len, 1) # We also test the unpacking. _, _, pred_y = self.model.evaluate(x, y, batch_size=ModelTest.batch_size, return_pred=True, callbacks=[self.mock_callback]) self.assertEqual(pred_y.shape, (ModelTest.evaluate_dataset_len, 1)) def test_evaluate_with_return_dict(self): x = torch.rand(ModelTest.evaluate_dataset_len, 1) y = torch.rand(ModelTest.evaluate_dataset_len, 1) logs = self.model.evaluate(x, y, batch_size=ModelTest.batch_size, return_dict_format=True) self._test_return_dict_logs(logs) def test_evaluate_with_np_array(self): x = np.random.rand(ModelTest.evaluate_dataset_len, 1).astype(np.float32) y = np.random.rand(ModelTest.evaluate_dataset_len, 1).astype(np.float32) loss, metrics, pred_y = self.model.evaluate( x, y, batch_size=ModelTest.batch_size, return_pred=True) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values + self.epoch_metrics_values) self.assertEqual(pred_y.shape, (ModelTest.evaluate_dataset_len, 1)) def test_evaluate_data_loader(self): x = torch.rand(ModelTest.evaluate_dataset_len, 1) y = torch.rand(ModelTest.evaluate_dataset_len, 1) dataset = TensorDataset(x, y) generator = DataLoader(dataset, ModelTest.batch_size) loss, metrics, pred_y = self.model.evaluate_generator(generator, return_pred=True) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values + self.epoch_metrics_values) self.assertEqual(pred_y.shape, (ModelTest.evaluate_dataset_len, 1)) def test_evaluate_generator(self): num_steps = 10 generator = some_data_tensor_generator(ModelTest.batch_size) loss, metrics, pred_y = self.model.evaluate_generator(generator, steps=num_steps, return_pred=True) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values + self.epoch_metrics_values) self.assertEqual(type(pred_y), np.ndarray) self.assertEqual(pred_y.shape, (num_steps * ModelTest.batch_size, 1)) def test_evaluate_generator_with_stop_iteration(self): test_generator = SomeDataGeneratorUsingStopIteration( ModelTest.batch_size, 10) loss, _ = self.model.evaluate_generator(test_generator) self.assertEqual(type(loss), float) def test_evaluate_generator_with_callback(self): num_steps = 10 generator = some_data_tensor_generator(ModelTest.batch_size) self.model.evaluate_generator(generator, steps=num_steps, callbacks=[self.mock_callback]) params = {'steps': ModelTest.epochs} self._test_callbacks_test(params) def test_evaluate_generator_with_return_dict(self): num_steps = 10 generator = some_data_tensor_generator(ModelTest.batch_size) logs = self.model.evaluate_generator(generator, steps=num_steps, return_dict_format=True) self._test_return_dict_logs(logs) def test_evaluate_generator_with_ground_truth(self): num_steps = 10 generator = some_data_tensor_generator(ModelTest.batch_size) loss, metrics, pred_y, true_y = self.model.evaluate_generator( generator, steps=num_steps, return_pred=True, return_ground_truth=True) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values + self.epoch_metrics_values) self.assertEqual(type(pred_y), np.ndarray) self.assertEqual(type(true_y), np.ndarray) self.assertEqual(pred_y.shape, (num_steps * ModelTest.batch_size, 1)) self.assertEqual(true_y.shape, (num_steps * ModelTest.batch_size, 1)) def test_evaluate_generator_with_no_concatenation(self): num_steps = 10 generator = some_data_tensor_generator(ModelTest.batch_size) loss, metrics, pred_y, true_y = self.model.evaluate_generator( generator, steps=num_steps, return_pred=True, return_ground_truth=True, concatenate_returns=False) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values + self.epoch_metrics_values) self.assertEqual(type(pred_y), list) for pred in pred_y: self.assertEqual(type(pred), np.ndarray) self.assertEqual(pred.shape, (ModelTest.batch_size, 1)) self.assertEqual(type(true_y), list) for true in true_y: self.assertEqual(type(true), np.ndarray) self.assertEqual(true.shape, (ModelTest.batch_size, 1)) def test_evaluate_with_only_one_metric(self): model = Model(self.pytorch_network, self.optimizer, self.loss_function, batch_metrics=self.batch_metrics[:1]) x = torch.rand(ModelTest.evaluate_dataset_len, 1) y = torch.rand(ModelTest.evaluate_dataset_len, 1) loss, first_metric = model.evaluate(x, y, batch_size=ModelTest.batch_size) self.assertEqual(type(loss), float) self.assertEqual(type(first_metric), float) self.assertEqual(first_metric, some_metric_1_value) def test_metrics_integration(self): num_steps = 10 model = Model(self.pytorch_network, self.optimizer, self.loss_function, batch_metrics=[F.mse_loss]) train_generator = some_data_tensor_generator(ModelTest.batch_size) valid_generator = some_data_tensor_generator(ModelTest.batch_size) model.fit_generator(train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch, callbacks=[self.mock_callback]) generator = some_data_tensor_generator(ModelTest.batch_size) loss, mse = model.evaluate_generator(generator, steps=num_steps) self.assertEqual(type(loss), float) self.assertEqual(type(mse), float) def test_epoch_metrics_integration(self): model = Model(self.pytorch_network, self.optimizer, self.loss_function, epoch_metrics=[SomeEpochMetric()]) train_generator = some_data_tensor_generator(ModelTest.batch_size) valid_generator = some_data_tensor_generator(ModelTest.batch_size) logs = model.fit_generator(train_generator, valid_generator, epochs=1, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch) actual_value = logs[-1]['some_epoch_metric'] val_actual_value = logs[-1]['val_some_epoch_metric'] expected_value = 5 self.assertEqual(val_actual_value, expected_value) self.assertEqual(actual_value, expected_value) def test_evaluate_with_no_metric(self): model = Model(self.pytorch_network, self.optimizer, self.loss_function) x = torch.rand(ModelTest.evaluate_dataset_len, 1) y = torch.rand(ModelTest.evaluate_dataset_len, 1) loss = model.evaluate(x, y, batch_size=ModelTest.batch_size) self.assertEqual(type(loss), float) def test_tensor_evaluate_on_batch(self): x = torch.rand(ModelTest.batch_size, 1) y = torch.rand(ModelTest.batch_size, 1) loss, metrics = self.model.evaluate_on_batch(x, y) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values) def test_evaluate_on_batch_with_pred(self): x = torch.rand(ModelTest.batch_size, 1) y = torch.rand(ModelTest.batch_size, 1) loss, metrics, pred_y = self.model.evaluate_on_batch(x, y, return_pred=True) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values) self.assertEqual(pred_y.shape, (ModelTest.batch_size, 1)) def test_ndarray_evaluate_on_batch(self): x = np.random.rand(ModelTest.batch_size, 1).astype(np.float32) y = np.random.rand(ModelTest.batch_size, 1).astype(np.float32) loss, metrics = self.model.evaluate_on_batch(x, y) self.assertEqual(type(loss), float) self.assertEqual(type(metrics), np.ndarray) self.assertEqual(metrics.tolist(), self.batch_metrics_values) def test_predict(self): x = torch.rand(ModelTest.evaluate_dataset_len, 1) pred_y = self.model.predict(x, batch_size=ModelTest.batch_size) self.assertEqual(pred_y.shape, (ModelTest.evaluate_dataset_len, 1)) def test_predict_with_np_array(self): x = np.random.rand(ModelTest.evaluate_dataset_len, 1).astype(np.float32) pred_y = self.model.predict(x, batch_size=ModelTest.batch_size) self.assertEqual(type(pred_y), np.ndarray) self.assertEqual(pred_y.shape, (ModelTest.evaluate_dataset_len, 1)) def test_predict_data_loader(self): x = torch.rand(ModelTest.evaluate_dataset_len, 1) generator = DataLoader(x, ModelTest.batch_size) pred_y = self.model.predict_generator(generator) self.assertEqual(type(pred_y), np.ndarray) self.assertEqual(pred_y.shape, (ModelTest.evaluate_dataset_len, 1)) def test_predict_generator(self): num_steps = 10 generator = some_data_tensor_generator(ModelTest.batch_size) generator = (x for x, _ in generator) pred_y = self.model.predict_generator(generator, steps=num_steps) self.assertEqual(type(pred_y), np.ndarray) self.assertEqual(pred_y.shape, (num_steps * ModelTest.batch_size, 1)) def test_predict_generator_with_no_concatenation(self): num_steps = 10 generator = some_data_tensor_generator(ModelTest.batch_size) generator = (x for x, _ in generator) pred_y = self.model.predict_generator(generator, steps=num_steps, concatenate_returns=False) self.assertEqual(type(pred_y), list) for pred in pred_y: self.assertEqual(type(pred), np.ndarray) self.assertEqual(pred.shape, (ModelTest.batch_size, 1)) def test_tensor_predict_on_batch(self): x = torch.rand(ModelTest.batch_size, 1) pred_y = self.model.predict_on_batch(x) self.assertEqual(pred_y.shape, (ModelTest.batch_size, 1)) def test_ndarray_predict_on_batch(self): x = np.random.rand(ModelTest.batch_size, 1).astype(np.float32) pred_y = self.model.predict_on_batch(x) self.assertEqual(pred_y.shape, (ModelTest.batch_size, 1)) @skipIf(not torch.cuda.is_available(), "no gpu available") def test_cpu_cuda(self): train_generator = some_data_tensor_generator(ModelTest.batch_size) valid_generator = some_data_tensor_generator(ModelTest.batch_size) self._capture_output() with torch.cuda.device(ModelTest.cuda_device): self.model.cuda() self.model.fit_generator( train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch, callbacks=[self.mock_callback]) # The context manager is also used here because of this bug: # https://github.com/pytorch/pytorch/issues/7320 with torch.cuda.device(ModelTest.cuda_device): self.model.cuda(ModelTest.cuda_device) self._test_device( torch.device('cuda:' + str(ModelTest.cuda_device))) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch, callbacks=[self.mock_callback]) self.model.cpu() self._test_device(torch.device('cpu')) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch, callbacks=[self.mock_callback]) self.model.to(torch.device('cuda:' + str(ModelTest.cuda_device))) self._test_device( torch.device('cuda:' + str(ModelTest.cuda_device))) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch, callbacks=[self.mock_callback]) self.model.to(torch.device('cpu')) self._test_device(torch.device('cpu')) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTest.epochs, steps_per_epoch=ModelTest.steps_per_epoch, validation_steps=ModelTest.steps_per_epoch, callbacks=[self.mock_callback]) def test_get_batch_size(self): batch_size = ModelTest.batch_size x = np.random.rand(batch_size, 1).astype(np.float32) y = np.random.rand(batch_size, 1).astype(np.float32) batch_size2 = ModelTest.batch_size + 1 x2 = np.random.rand(batch_size2, 1).astype(np.float32) y2 = np.random.rand(batch_size2, 1).astype(np.float32) other_batch_size = batch_size2 + 1 inf_batch_size = self.model.get_batch_size(x, y) self.assertEqual(inf_batch_size, batch_size) inf_batch_size = self.model.get_batch_size(x2, y2) self.assertEqual(inf_batch_size, batch_size2) inf_batch_size = self.model.get_batch_size(x, y2) self.assertEqual(inf_batch_size, batch_size) inf_batch_size = self.model.get_batch_size(x2, y) self.assertEqual(inf_batch_size, batch_size2) inf_batch_size = self.model.get_batch_size((x, x2), y) self.assertEqual(inf_batch_size, batch_size) inf_batch_size = self.model.get_batch_size((x2, x), y) self.assertEqual(inf_batch_size, batch_size) inf_batch_size = self.model.get_batch_size((x, x2), (y, y2)) self.assertEqual(inf_batch_size, batch_size) inf_batch_size = self.model.get_batch_size((x2, x), (y, y2)) self.assertEqual(inf_batch_size, batch_size2) inf_batch_size = self.model.get_batch_size([x, x2], y) self.assertEqual(inf_batch_size, batch_size) inf_batch_size = self.model.get_batch_size([x2, x], y) self.assertEqual(inf_batch_size, batch_size) inf_batch_size = self.model.get_batch_size([x, x2], [y, y2]) self.assertEqual(inf_batch_size, batch_size) inf_batch_size = self.model.get_batch_size([x2, x], [y, y2]) self.assertEqual(inf_batch_size, batch_size2) inf_batch_size = self.model.get_batch_size( { 'batch_size': other_batch_size, 'x': x }, {'y': y}) self.assertEqual(inf_batch_size, other_batch_size) inf_batch_size = self.model.get_batch_size({'x': x}, { 'batch_size': other_batch_size, 'y': y }) self.assertEqual(inf_batch_size, other_batch_size) inf_batch_size = self.model.get_batch_size({'x': x}, {'y': y}) self.assertEqual(inf_batch_size, batch_size) inf_batch_size = self.model.get_batch_size( OrderedDict([('x1', x), ('x2', x2)]), {'y': y}) self.assertEqual(inf_batch_size, batch_size) inf_batch_size = self.model.get_batch_size( OrderedDict([('x1', x2), ('x2', x)]), {'y': y}) self.assertEqual(inf_batch_size, batch_size2) inf_batch_size = self.model.get_batch_size([1, 2, 3], {'y': y}) self.assertEqual(inf_batch_size, batch_size) def test_get_batch_size_warning(self): with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") inf_batch_size = self.model.get_batch_size([1, 2, 3], [4, 5, 6]) self.assertEqual(inf_batch_size, 1) self.assertEqual(len(w), 1) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") warning_settings['batch_size'] = 'ignore' inf_batch_size = self.model.get_batch_size([1, 2, 3], [4, 5, 6]) self.assertEqual(inf_batch_size, 1) self.assertEqual(len(w), 0)
class BaseCSVGradientLoggerTest: # pylint: disable=not-callable,no-member CSVGradientLogger = None batch_size = 20 num_epochs = 10 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Sequential(nn.Linear(1, 2), nn.Linear(2, 1)) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) self.temp_dir_obj = TemporaryDirectory() self.csv_filename = os.path.join(self.temp_dir_obj.name, 'layer_{}.csv') def tearDown(self): self.temp_dir_obj.cleanup() def test_logging(self): train_gen = some_data_generator(self.batch_size) valid_gen = some_data_generator(self.batch_size) memgrad = MemoryGradientLogger() logger = self.CSVGradientLogger(self.csv_filename) self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[memgrad, logger]) self._test_logging(memgrad.history) def test_logging_append(self): train_gen = some_data_generator(self.batch_size) valid_gen = some_data_generator(self.batch_size) logger = self.CSVGradientLogger(self.csv_filename) memgrad = MemoryGradientLogger() self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[memgrad, logger]) memgrad2 = MemoryGradientLogger() logger = self.CSVGradientLogger(self.csv_filename, append=True) self.model.fit_generator( train_gen, valid_gen, epochs=20, steps_per_epoch=5, initial_epoch=self.num_epochs, callbacks=[memgrad2, logger], ) history = { layer: stats1 + memgrad2.history[layer] for layer, stats1 in memgrad.history.items() } self._test_logging(history) def test_logging_overwrite(self): train_gen = some_data_generator(self.batch_size) valid_gen = some_data_generator(self.batch_size) logger = self.CSVGradientLogger(self.csv_filename) self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[logger]) memgrad = MemoryGradientLogger() logger = self.CSVGradientLogger(self.csv_filename, append=False) self.model.fit_generator( train_gen, valid_gen, epochs=20, steps_per_epoch=5, initial_epoch=self.num_epochs, callbacks=[memgrad, logger], ) self._test_logging(memgrad.history) def _test_logging(self, history): for layer, stats in history.items(): filename = self.csv_filename.format(layer) with open(filename, 'r', encoding='utf-8') as csvfile: reader = csv.DictReader(csvfile) rows = list(reader) self.assertEqual(len(rows), len(stats)) for row, stats_entry in zip(rows, stats): self.assertEqual(row.keys(), stats_entry.keys()) for k in row.keys(): self.assertAlmostEqual(float(row[k]), stats_entry[k])
from poutyne import Model model = Model(network, 'sgd', 'cross_entropy', batch_metrics=['accuracy'], epoch_metrics=['f1']) model.to(device) model.fit_generator(train_loader, valid_loader, epochs=num_epochs, callbacks=callbacks) test_loss, (test_acc, test_f1) = model.evaluate_generator(test_loader) print(f'Test: Loss: {test_loss}, Accuracy: {test_acc}, F1: {test_f1}')
class ModelTestMultiGPU(ModelFittingTestCase): # pylint: disable=too-many-public-methods def setUp(self): super().setUp() torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.batch_metrics = [ some_batch_metric_1, ('custom_name', some_batch_metric_2), repeat_batch_metric, repeat_batch_metric ] self.batch_metrics_names = [ 'some_batch_metric_1', 'custom_name', 'repeat_batch_metric1', 'repeat_batch_metric2' ] self.batch_metrics_values = [ some_metric_1_value, some_metric_2_value, repeat_batch_metric_value, repeat_batch_metric_value ] self.epoch_metrics = [SomeConstantEpochMetric()] self.epoch_metrics_names = ['some_constant_epoch_metric'] self.epoch_metrics_values = [some_constant_epoch_metric_value] self.model = Model(self.pytorch_network, self.optimizer, self.loss_function, batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics) self.default_main_device = ModelTestMultiGPU.cuda_device def _test_multiple_gpu_mode(self, devices): if devices == "all": expected = torch.cuda.device_count() else: expected = len(devices) self.assertEqual(len([self.model.device] + self.model.other_device), expected) def _test_single_gpu_mode(self): self.assertIsNone(self.model.other_device) self.assertEqual(len([self.model.device]), 1) def test_back_and_forth_gpu_cpu_multi_gpus(self): devices = "all" train_generator = some_data_tensor_generator( ModelTestMultiGPU.batch_size) valid_generator = some_data_tensor_generator( ModelTestMultiGPU.batch_size) with torch.cuda.device(self.default_main_device): self.model.to(devices) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_device( torch.device('cuda:' + str(self.default_main_device))) self._test_multiple_gpu_mode(devices=devices) self.model.cpu() self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_device(torch.device('cpu')) self._test_single_gpu_mode() self.model.to(devices) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_device( torch.device('cuda:' + str(ModelTestMultiGPU.cuda_device))) self._test_multiple_gpu_mode(devices=devices) self.model.to(torch.device('cpu')) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_device(torch.device('cpu')) self._test_single_gpu_mode() def test_back_and_forth_cuda_cpu_to_multi_gpus(self): devices = "all" train_generator = some_data_tensor_generator( ModelTestMultiGPU.batch_size) valid_generator = some_data_tensor_generator( ModelTestMultiGPU.batch_size) self.model.to(devices) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) with torch.cuda.device(self.default_main_device): self.model.cuda() self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_device( torch.device('cuda:' + str(self.default_main_device))) self._test_single_gpu_mode() self.model.cpu() self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_device(torch.device('cpu')) self._test_single_gpu_mode() self.model.to(devices) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_device( torch.device('cuda:' + str(ModelTestMultiGPU.cuda_device))) self._test_multiple_gpu_mode(devices=devices) self.model.cuda() self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_device( torch.device('cuda:' + str(self.default_main_device))) self._test_single_gpu_mode() self.model.to(torch.device('cpu')) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_device(torch.device('cpu')) self._test_single_gpu_mode() def test_devices_settings(self): train_generator = some_data_tensor_generator( ModelTestMultiGPU.batch_size) valid_generator = some_data_tensor_generator( ModelTestMultiGPU.batch_size) devices = "all" self.model.to(devices) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_multiple_gpu_mode(devices=devices) devices = ["cuda:0", "cuda:1"] self.model.to(devices) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_multiple_gpu_mode(devices=devices) devices = [torch.device("cuda:0"), torch.device("cuda:1")] self.model.to(devices) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_multiple_gpu_mode(devices=devices) devices = ["cuda:0"] self.model.to(devices) self.assertIsNone(self.model.other_device) self.model.fit_generator( train_generator, valid_generator, epochs=ModelTestMultiGPU.epochs, steps_per_epoch=ModelTestMultiGPU.steps_per_epoch, validation_steps=ModelTestMultiGPU.steps_per_epoch, callbacks=[self.mock_callback]) self._test_device(torch.device('cuda:0')) self._test_single_gpu_mode()
class ModelFittingTestCaseProgress(ModelFittingTestCase): # pylint: disable=too-many-public-methods num_steps = 5 def setUp(self): super().setUp() self.train_generator = some_data_tensor_generator(ModelFittingTestCase.batch_size) self.valid_generator = some_data_tensor_generator(ModelFittingTestCase.batch_size) self.test_generator = some_data_tensor_generator(ModelFittingTestCase.batch_size) torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.batch_metrics = [ some_batch_metric_1, ('custom_name', some_batch_metric_2), repeat_batch_metric, repeat_batch_metric ] self.batch_metrics_names = [ 'some_batch_metric_1', 'custom_name', 'repeat_batch_metric1', 'repeat_batch_metric2' ] self.batch_metrics_values = [ some_metric_1_value, some_metric_2_value, repeat_batch_metric_value, repeat_batch_metric_value ] self.epoch_metrics = [SomeConstantEpochMetric()] self.epoch_metrics_names = ['some_constant_epoch_metric'] self.epoch_metrics_values = [some_constant_epoch_metric_value] self.model = Model(self.pytorch_network, self.optimizer, self.loss_function, batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics) self._capture_output() def assertStdoutContains(self, values): for value in values: self.assertIn(value, self.test_out.getvalue().strip()) def assertStdoutNotContains(self, values): for value in values: self.assertNotIn(value, self.test_out.getvalue().strip()) @skipIf(color is None, "Unable to import colorama") def test_fitting_with_default_coloring(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback]) self.assertStdoutContains(["[32m", "[35m", "[36m", "[94m"]) def test_fitting_with_progress_bar_show_epoch(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback]) self.assertStdoutContains(["Epoch", "1/5", "2/5"]) def test_fitting_with_progress_bar_show_steps(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback]) self.assertStdoutContains(["steps", f"{ModelFittingTestCase.steps_per_epoch}"]) def test_fitting_with_progress_bar_show_train_val_final_steps(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback]) self.assertStdoutContains(["Val steps", "Train steps"]) def test_fitting_with_no_progress_bar_dont_show_epoch(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback], verbose=False) self.assertStdoutNotContains(["Epoch", "1/5", "2/5"]) @skipIf(color is None, "Unable to import colorama") def test_fitting_with_user_coloring(self): coloring = { "text_color": 'BLACK', "ratio_color": "BLACK", "metric_value_color": "BLACK", "time_color": "BLACK", "progress_bar_color": "BLACK" } _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback], progress_options=dict(coloring=coloring)) self.assertStdoutContains(["[30m"]) @skipIf(color is None, "Unable to import colorama") def test_fitting_with_user_partial_coloring(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback], progress_options=dict(coloring={ "text_color": 'BLACK', "ratio_color": "BLACK" })) self.assertStdoutContains(["[30m", "[32m", "[35m", "[94m"]) def test_fitting_with_user_coloring_invalid(self): with self.assertRaises(KeyError): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback], progress_options=dict(coloring={"invalid_name": 'A COLOR'})) def test_fitting_with_no_coloring(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback], progress_options=dict(coloring=False)) self.assertStdoutNotContains(["[32m", "[35m", "[36m", "[94m"]) @skipIf(color is None, "Unable to import colorama") def test_fitting_with_progress_bar_default_color(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback], progress_options=dict(coloring=True, progress_bar=True)) self.assertStdoutContains(["%", "[32m", "[35m", "[36m", "[94m", "\u2588"]) @skipIf(color is None, "Unable to import colorama") def test_fitting_with_progress_bar_user_color(self): coloring = { "text_color": 'BLACK', "ratio_color": "BLACK", "metric_value_color": "BLACK", "time_color": "BLACK", "progress_bar_color": "BLACK" } _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback], progress_options=dict(coloring=coloring, progress_bar=True)) self.assertStdoutContains(["%", "[30m", "\u2588"]) def test_fitting_with_progress_bar_no_color(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback], progress_options=dict(coloring=False, progress_bar=True)) self.assertStdoutContains(["%", "\u2588"]) self.assertStdoutNotContains(["[32m", "[35m", "[36m", "[94m"]) def test_fitting_with_no_progress_bar(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=ModelFittingTestCase.epochs, steps_per_epoch=ModelFittingTestCase.steps_per_epoch, validation_steps=ModelFittingTestCase.steps_per_epoch, callbacks=[self.mock_callback], progress_options=dict(coloring=False, progress_bar=False)) self.assertStdoutNotContains(["%", "\u2588"]) self.assertStdoutNotContains(["[32m", "[35m", "[36m", "[94m"]) def test_progress_bar_with_step_is_none(self): train_generator = SomeDataGeneratorUsingStopIteration(ModelFittingTestCase.batch_size, 10) valid_generator = SomeDataGeneratorUsingStopIteration(ModelFittingTestCase.batch_size, 10) _ = self.model.fit_generator(train_generator, valid_generator, epochs=ModelFittingTestCase.epochs, progress_options=dict(coloring=False, progress_bar=True)) self.assertStdoutContains(["s/step"]) self.assertStdoutNotContains(["[32m", "[35m", "[36m", "[94m", "\u2588", "%"]) def test_evaluate_without_progress_output(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) _, _ = self.model.evaluate(x, y, batch_size=ModelFittingTestCase.batch_size, verbose=False) self.assertStdoutNotContains(["[32m", "[35m", "[36m", "[94m"]) @skipIf(color is None, "Unable to import colorama") def test_evaluate_with_default_coloring(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) _, _ = self.model.evaluate(x, y, batch_size=ModelFittingTestCase.batch_size) self.assertStdoutContains(["[32m", "[35m", "[36m", "[94m"]) @skipIf(color is None, "Unable to import colorama") def test_evaluate_with_user_coloring(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) coloring = { "text_color": 'BLACK', "ratio_color": "BLACK", "metric_value_color": "BLACK", "time_color": "BLACK", "progress_bar_color": "BLACK" } _, _ = self.model.evaluate(x, y, batch_size=ModelFittingTestCase.batch_size, progress_options=dict(coloring=coloring)) self.assertStdoutContains(["[30m"]) @skipIf(color is None, "Unable to import colorama") def test_evaluate_with_user_partial_coloring(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) _, _ = self.model.evaluate(x, y, batch_size=ModelFittingTestCase.batch_size, progress_options=dict(coloring={ "text_color": 'BLACK', "ratio_color": "BLACK" })) self.assertStdoutContains(["[30m", "[32m", "[35m", "[94m"]) def test_evaluate_with_user_coloring_invalid(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) with self.assertRaises(KeyError): _, _ = self.model.evaluate(x, y, batch_size=ModelFittingTestCase.batch_size, progress_options=dict(coloring={"invalid_name": 'A COLOR'})) def test_evaluate_with_no_coloring(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) _, _ = self.model.evaluate(x, y, batch_size=ModelFittingTestCase.batch_size, progress_options=dict(coloring=False)) self.assertStdoutNotContains(["[32m", "[35m", "[36m", "[94m"]) @skipIf(color is None, "Unable to import colorama") def test_evaluate_with_progress_bar_default_coloring(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) _, _ = self.model.evaluate(x, y, batch_size=ModelFittingTestCase.batch_size, progress_options=dict(coloring=True, progress_bar=True)) self.assertStdoutContains(["%", "[32m", "[35m", "[36m", "[94m", "\u2588"]) @skipIf(color is None, "Unable to import colorama") def test_evaluate_with_progress_bar_user_coloring(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) coloring = { "text_color": 'BLACK', "ratio_color": "BLACK", "metric_value_color": "BLACK", "time_color": "BLACK", "progress_bar_color": "BLACK" } _, _ = self.model.evaluate(x, y, batch_size=ModelFittingTestCase.batch_size, progress_options=dict(coloring=coloring, progress_bar=True)) self.assertStdoutContains(["%", "[30m", "\u2588"]) @skipIf(color is None, "Unable to import colorama") def test_evaluate_with_progress_bar_user_no_color(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) _, _ = self.model.evaluate(x, y, batch_size=ModelFittingTestCase.batch_size, progress_options=dict(coloring=False, progress_bar=True)) self.assertStdoutContains(["%", "\u2588"]) self.assertStdoutNotContains(["[32m", "[35m", "[36m", "[94m"]) def test_evaluate_with_no_progress_bar(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) _, _ = self.model.evaluate(x, y, batch_size=ModelFittingTestCase.batch_size, progress_options=dict(coloring=False, progress_bar=False)) self.assertStdoutNotContains(["%", "\u2588"]) self.assertStdoutNotContains(["[32m", "[35m", "[36m", "[94m"]) def test_evaluate_data_loader_with_progress_bar_coloring(self): x = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) y = torch.rand(ModelFittingTestCase.evaluate_dataset_len, 1) dataset = TensorDataset(x, y) generator = DataLoader(dataset, ModelFittingTestCase.batch_size) _, _ = self.model.evaluate_generator(generator, verbose=True) self.assertStdoutContains(["%", "[32m", "[35m", "[36m", "[94m", "\u2588"]) def test_evaluate_generator_with_progress_bar_coloring(self): generator = some_data_tensor_generator(ModelFittingTestCase.batch_size) _, _ = self.model.evaluate_generator(generator, steps=ModelFittingTestCaseProgress.num_steps, verbose=True) self.assertStdoutContains(["%", "[32m", "[35m", "[36m", "[94m", "\u2588"]) def test_evaluate_generator_with_callback_and_progress_bar_coloring(self): generator = some_data_tensor_generator(ModelFittingTestCase.batch_size) _, _ = self.model.evaluate_generator(generator, steps=ModelFittingTestCaseProgress.num_steps, callbacks=[self.mock_callback], verbose=True) self.assertStdoutContains(["%", "[32m", "[35m", "[36m", "[94m", "\u2588"]) def test_fitting_complete_display_test_with_progress_bar_coloring(self): # we use the same color for all components for simplicity coloring = { "text_color": 'WHITE', "ratio_color": "WHITE", "metric_value_color": "WHITE", "time_color": "WHITE", "progress_bar_color": "WHITE" } _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=1, steps_per_epoch=ModelFittingTestCaseProgress.num_steps, validation_steps=ModelFittingTestCaseProgress.num_steps, callbacks=[self.mock_callback], progress_options=dict(coloring=coloring, progress_bar=False)) # We split per step update steps_update = self.test_out.getvalue().strip().split("\r") # we don't validate the templating of metrics since tested before template_format = r".*Epoch:.*{}\/1.*\[37mStep:.*{}\/5.*{:6.2f}\%.*|{}|.*ETA:" epoch = 1 # the 5 train steps for step, step_update in enumerate(steps_update[:ModelFittingTestCaseProgress.num_steps]): step += 1 progress_Bar = "\u2588" * step * 2 + " " * (20 - step * 2) regex_filled = template_format.format(epoch, step, step / ModelFittingTestCaseProgress.num_steps * 100, progress_Bar) self.assertRegex(step_update, regex_filled) # The 5 val steps for step, step_update in enumerate(steps_update[ModelFittingTestCaseProgress.num_steps:-1]): step += 1 progress_Bar = "\u2588" * step * 2 + " " * (20 - step * 2) regex_filled = template_format.format(epoch, step, step / ModelFittingTestCaseProgress.num_steps * 100, progress_Bar) self.assertRegex(step_update, regex_filled) # last print update templating different last_print_regex = r".*\[37mTrain steps:.*5.*Val steps:.*5.*[0-9]*\.[0-9][0-9]s" self.assertRegex(steps_update[-1], last_print_regex) def test_fitting_complete_display_test_with_progress_bar_no_coloring(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=1, steps_per_epoch=ModelFittingTestCaseProgress.num_steps, validation_steps=ModelFittingTestCaseProgress.num_steps, callbacks=[self.mock_callback], progress_options=dict(coloring=False, progress_bar=True)) # We split per step update steps_update = self.test_out.getvalue().strip().split("\r") # we don't validate the templating of metrics since tested before template_format = r".*Epoch:.*{}\/1.*Step:.*{}\/5.*{:6.2f}\%.*|{}|.*ETA:" epoch = 1 # the 5 train steps for step, step_update in enumerate(steps_update[:ModelFittingTestCaseProgress.num_steps]): step += 1 progress_Bar = "\u2588" * step * 2 + " " * (20 - step * 2) regex_filled = template_format.format(epoch, step, step / ModelFittingTestCaseProgress.num_steps * 100, progress_Bar) self.assertRegex(step_update, regex_filled) # The 5 val steps for step, step_update in enumerate(steps_update[ModelFittingTestCaseProgress.num_steps:-1]): step += 1 progress_Bar = "\u2588" * step * 2 + " " * (20 - step * 2) regex_filled = template_format.format(epoch, step, step / ModelFittingTestCaseProgress.num_steps * 100, progress_Bar) self.assertRegex(step_update, regex_filled) # last print update templating different last_print_regex = r".*Train steps:.*5.*Val steps:.*5.*[0-9]*\.[0-9][0-9]s" self.assertRegex(steps_update[-1], last_print_regex) def test_fitting_complete_display_test_with_no_progress_bar_no_coloring(self): _ = self.model.fit_generator(self.train_generator, self.valid_generator, epochs=1, steps_per_epoch=ModelFittingTestCaseProgress.num_steps, validation_steps=ModelFittingTestCaseProgress.num_steps, callbacks=[self.mock_callback], progress_options=dict(coloring=False, progress_bar=False)) # We split per step update steps_update = self.test_out.getvalue().strip().split("\r") # we don't validate the templating of metrics since tested before template_format = r".*Epoch:.*{}\/1.*Step:.*{}\/5.*ETA:" epoch = 1 # the 5 train steps for step, step_update in enumerate(steps_update[:ModelFittingTestCaseProgress.num_steps]): step += 1 regex_filled = template_format.format(epoch, step, step / ModelFittingTestCaseProgress.num_steps * 100) self.assertRegex(step_update, regex_filled) # The 5 val steps for step, step_update in enumerate(steps_update[ModelFittingTestCaseProgress.num_steps:-1]): step += 1 regex_filled = template_format.format(epoch, step, step / ModelFittingTestCaseProgress.num_steps * 100) self.assertRegex(step_update, regex_filled) # last print update templating different last_print_regex = r".*Train steps:.*5.*Val steps:.*5.*[0-9]*\.[0-9][0-9]s" self.assertRegex(steps_update[-1], last_print_regex) def test_evaluate_complete_display_test_with_progress_bar_coloring(self): # we use the same color for all components for simplicity coloring = { "text_color": 'WHITE', "ratio_color": "WHITE", "metric_value_color": "WHITE", "time_color": "WHITE", "progress_bar_color": "WHITE" } _, _ = self.model.evaluate_generator(self.test_generator, steps=ModelFittingTestCaseProgress.num_steps, callbacks=[self.mock_callback], verbose=True, progress_options=dict(coloring=coloring, progress_bar=True)) # We split per step update steps_update = self.test_out.getvalue().strip().split("\r") # we don't validate the templating of metrics since tested before template_format = r".*\[37mStep:.*{}\/5.*{:6.2f}\%.*|{}|.*ETA:" for step, step_update in enumerate(steps_update[:-1]): step += 1 progress_Bar = "\u2588" * step * 2 + " " * (20 - step * 2) regex_filled = template_format.format(step, step / ModelFittingTestCaseProgress.num_steps * 100, progress_Bar) self.assertRegex(step_update, regex_filled) # last print update templating different last_print_regex = r".*\[37mTest steps:.*5.*[0-9]*\.[0-9][0-9]s" self.assertRegex(steps_update[-1], last_print_regex) def test_evaluate_complete_display_test_with_progress_bar_no_coloring(self): _, _ = self.model.evaluate_generator(self.test_generator, steps=ModelFittingTestCaseProgress.num_steps, callbacks=[self.mock_callback], verbose=True, progress_options=dict(coloring=False, progress_bar=True)) # We split per step update steps_update = self.test_out.getvalue().strip().split("\r") # we don't validate the templating of metrics since tested before template_format = r".*Step:.*{}\/5.*{:6.2f}\%.*|{}|.*ETA:" for step, step_update in enumerate(steps_update[:-1]): step += 1 progress_Bar = "\u2588" * step * 2 + " " * (20 - step * 2) regex_filled = template_format.format(step, step / ModelFittingTestCaseProgress.num_steps * 100, progress_Bar) self.assertRegex(step_update, regex_filled) # last print update templating different last_print_regex = r".*Test steps:.*5.*[0-9]*\.[0-9][0-9]s" self.assertRegex(steps_update[-1], last_print_regex) def test_evaluate_complete_display_test_with_no_progress_bar_no_coloring(self): _, _ = self.model.evaluate_generator(self.test_generator, steps=ModelFittingTestCaseProgress.num_steps, callbacks=[self.mock_callback], verbose=True, progress_options=dict(coloring=False, progress_bar=False)) # We split per step update steps_update = self.test_out.getvalue().strip().split("\r") # we don't validate the templating of metrics since tested before template_format = r".*Step:.*{}\/5.*ETA:" for step, step_update in enumerate(steps_update[:-1]): step += 1 regex_filled = template_format.format(step, step / ModelFittingTestCaseProgress.num_steps * 100) self.assertRegex(step_update, regex_filled) # last print update templating different last_print_regex = r".*Test steps:.*5.*[0-9]*\.[0-9][0-9]s" self.assertRegex(steps_update[-1], last_print_regex)
class ModelCheckpointTest(TestCase): batch_size = 20 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) self.temp_dir_obj = TemporaryDirectory() self.checkpoint_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint_{epoch}.ckpt') def tearDown(self): self.temp_dir_obj.cleanup() def test_integration(self): train_gen = some_data_generator(ModelCheckpointTest.batch_size) valid_gen = some_data_generator(ModelCheckpointTest.batch_size) checkpointer = ModelCheckpoint(self.checkpoint_filename, monitor='val_loss', verbose=True, save_best_only=True) self.model.fit_generator(train_gen, valid_gen, epochs=10, steps_per_epoch=5, callbacks=[checkpointer]) def test_temporary_filename_arg(self): tmp_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint.tmp.ckpt') checkpoint_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint.ckpt') train_gen = some_data_generator(ModelCheckpointTest.batch_size) valid_gen = some_data_generator(ModelCheckpointTest.batch_size) checkpointer = ModelCheckpoint(checkpoint_filename, monitor='val_loss', verbose=True, period=1, temporary_filename=tmp_filename) self.model.fit_generator(train_gen, valid_gen, epochs=10, steps_per_epoch=5, callbacks=[checkpointer]) self.assertFalse(os.path.isfile(tmp_filename)) self.assertTrue(os.path.isfile(checkpoint_filename)) def test_temporary_filename_arg_with_differing_checkpoint_filename(self): epochs = 10 tmp_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint.tmp.ckpt') checkpoint_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint_{epoch}.ckpt') train_gen = some_data_generator(ModelCheckpointTest.batch_size) valid_gen = some_data_generator(ModelCheckpointTest.batch_size) checkpointer = ModelCheckpoint(checkpoint_filename, monitor='val_loss', verbose=True, period=1, temporary_filename=tmp_filename) self.model.fit_generator(train_gen, valid_gen, epochs=epochs, steps_per_epoch=5, callbacks=[checkpointer]) self.assertFalse(os.path.isfile(tmp_filename)) for i in range(1, epochs + 1): self.assertTrue(os.path.isfile( checkpoint_filename.format(epoch=i))) def test_non_atomic_write(self): checkpoint_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint.ckpt') train_gen = some_data_generator(ModelCheckpointTest.batch_size) valid_gen = some_data_generator(ModelCheckpointTest.batch_size) checkpointer = ModelCheckpoint(checkpoint_filename, monitor='val_loss', verbose=True, period=1, atomic_write=False) self.model.fit_generator(train_gen, valid_gen, epochs=10, steps_per_epoch=5, callbacks=[checkpointer]) self.assertTrue(os.path.isfile(checkpoint_filename)) def test_save_best_only(self): checkpointer = ModelCheckpoint(self.checkpoint_filename, monitor='val_loss', verbose=True, save_best_only=True) val_losses = [10, 3, 8, 5, 2] has_checkpoints = [True, True, False, False, True] self._test_checkpointer_with_val_losses(checkpointer, val_losses, has_checkpoints) def test_save_best_only_with_restore_best(self): checkpointer = ModelCheckpoint(self.checkpoint_filename, monitor='val_loss', verbose=True, save_best_only=True, restore_best=True) val_losses = [10, 3, 8, 5, 2] has_checkpoints = [True, True, False, False, True] self._test_checkpointer_with_val_losses(checkpointer, val_losses, has_checkpoints) self._test_restore_best(val_losses) def test_restore_best_without_save_best_only(self): with self.assertRaises(ValueError): ModelCheckpoint(self.checkpoint_filename, monitor='val_loss', verbose=True, save_best_only=False, restore_best=True) with self.assertRaises(ValueError): ModelCheckpoint(self.checkpoint_filename, monitor='val_loss', verbose=True, restore_best=True) def test_save_best_only_with_max(self): checkpointer = ModelCheckpoint(self.checkpoint_filename, monitor='val_loss', mode='max', verbose=True, save_best_only=True) val_losses = [2, 3, 8, 5, 2] has_checkpoints = [True, True, True, False, False] self._test_checkpointer_with_val_losses(checkpointer, val_losses, has_checkpoints) def test_periodic_with_period_of_1(self): checkpointer = ModelCheckpoint(self.checkpoint_filename, monitor='val_loss', verbose=True, period=1, save_best_only=False) val_losses = [1] * 10 has_checkpoints = [True] * 10 self._test_checkpointer_with_val_losses(checkpointer, val_losses, has_checkpoints) def test_periodic_with_period_of_2(self): checkpointer = ModelCheckpoint(self.checkpoint_filename, monitor='val_loss', verbose=True, period=2, save_best_only=False) val_losses = [1] * 10 has_checkpoints = [False, True] * 5 self._test_checkpointer_with_val_losses(checkpointer, val_losses, has_checkpoints) def _test_checkpointer_with_val_losses(self, checkpointer, val_losses, has_checkpoints): generator = some_data_generator(ModelCheckpointTest.batch_size) checkpointer.set_params({'epochs': len(val_losses), 'steps': 1}) checkpointer.set_model(self.model) checkpointer.on_train_begin({}) for epoch, (val_loss, has_checkpoint) in enumerate( zip(val_losses, has_checkpoints), 1): checkpointer.on_epoch_begin(epoch, {}) checkpointer.on_train_batch_begin(1, {}) loss = self._update_model(generator) checkpointer.on_train_batch_end(1, { 'batch': 1, 'size': ModelCheckpointTest.batch_size, 'loss': loss }) checkpointer.on_epoch_end(epoch, { 'epoch': epoch, 'loss': loss, 'val_loss': val_loss }) filename = self.checkpoint_filename.format(epoch=epoch) self.assertEqual(has_checkpoint, os.path.isfile(filename)) checkpointer.on_train_end({}) def _update_model(self, generator): self.pytorch_network.zero_grad() x, y = next(generator) pred_y = self.pytorch_network(x) loss = self.loss_function(pred_y, y) loss.backward() self.optimizer.step() return float(loss) def _test_restore_best(self, val_losses): final_weights = torch_to_numpy(self.model.get_weight_copies()) epoch = val_losses.index(min(val_losses)) + 1 best_epoch_filename = self.checkpoint_filename.format(epoch=epoch) self.model.load_weights(best_epoch_filename) best_weights = torch_to_numpy(self.model.get_weight_copies()) self.assertEqual(best_weights, final_weights)
class OptimizerCheckpointTest(TestCase): batch_size = 20 epochs = 10 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.Adam(self.pytorch_network.parameters(), lr=1e-3) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) self.temp_dir_obj = TemporaryDirectory() self.checkpoint_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint_{epoch}.optim') def tearDown(self): self.temp_dir_obj.cleanup() def test_integration(self): train_gen = some_data_generator(OptimizerCheckpointTest.batch_size) valid_gen = some_data_generator(OptimizerCheckpointTest.batch_size) checkpointer = OptimizerCheckpoint(self.checkpoint_filename, period=1) self.model.fit_generator(train_gen, valid_gen, epochs=OptimizerCheckpointTest.epochs, steps_per_epoch=5, callbacks=[checkpointer]) def test_checkpoints(self): checkpointer = OptimizerCheckpoint(self.checkpoint_filename, period=1) self._test_checkpointer(checkpointer) def _test_checkpointer(self, checkpointer): optimizer_states = {} generator = some_data_generator(OptimizerCheckpointTest.batch_size) checkpointer.set_params({ 'epochs': OptimizerCheckpointTest.epochs, 'steps': 1 }) checkpointer.set_model(self.model) checkpointer.on_train_begin({}) for epoch in range(1, OptimizerCheckpointTest.epochs + 1): checkpointer.on_epoch_begin(epoch, {}) checkpointer.on_train_batch_begin(1, {}) loss = self._update_model(generator) checkpointer.on_train_batch_end( 1, { 'batch': 1, 'size': OptimizerCheckpointTest.batch_size, 'loss': loss }) checkpointer.on_epoch_end(epoch, { 'epoch': epoch, 'loss': loss, 'val_loss': 1 }) filename = self.checkpoint_filename.format(epoch=epoch) self.assertTrue(os.path.isfile(filename)) optimizer_states[epoch] = torch_to_numpy( self.optimizer.state_dict(), copy=True) checkpointer.on_train_end({}) self._test_checkpoint(optimizer_states) def _update_model(self, generator): self.pytorch_network.zero_grad() x, y = next(generator) pred_y = self.pytorch_network(x) loss = self.loss_function(pred_y, y) loss.backward() self.optimizer.step() return float(loss) def _test_checkpoint(self, optimizer_states): for epoch, epoch_optimizer_state in optimizer_states.items(): filename = self.checkpoint_filename.format(epoch=epoch) self.model.load_optimizer_state(filename) saved_optimizer_state = torch_to_numpy(self.optimizer.state_dict()) self.assertEqual(epoch_optimizer_state, saved_optimizer_state)
class BaseCSVLoggerTest: # pylint: disable=not-callable,no-member CSVLogger = None batch_size = 20 lr = 1e-3 num_epochs = 10 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=BaseCSVLoggerTest.lr) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) self.temp_dir_obj = TemporaryDirectory() self.csv_filename = os.path.join(self.temp_dir_obj.name, 'my_log.csv') def tearDown(self): self.temp_dir_obj.cleanup() def test_logging(self): train_gen = some_data_generator(20) valid_gen = some_data_generator(20) logger = self.CSVLogger(self.csv_filename) history = self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[logger]) self._test_logging(history) def test_logging_with_batch_granularity(self): train_gen = some_data_generator(20) valid_gen = some_data_generator(20) logger = self.CSVLogger(self.csv_filename, batch_granularity=True) history = History() self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[logger, history]) self._test_logging(history.history) def test_logging_append(self): train_gen = some_data_generator(20) valid_gen = some_data_generator(20) logger = self.CSVLogger(self.csv_filename) history = self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[logger]) logger = self.CSVLogger(self.csv_filename, append=True) history2 = self.model.fit_generator(train_gen, valid_gen, epochs=20, steps_per_epoch=5, initial_epoch=self.num_epochs, callbacks=[logger]) self._test_logging(history + history2) def test_logging_overwrite(self): train_gen = some_data_generator(20) valid_gen = some_data_generator(20) logger = self.CSVLogger(self.csv_filename) self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[logger]) logger = self.CSVLogger(self.csv_filename, append=False) history = self.model.fit_generator(train_gen, valid_gen, epochs=20, steps_per_epoch=5, initial_epoch=self.num_epochs, callbacks=[logger]) self._test_logging(history) def _test_logging(self, history): with open(self.csv_filename) as csvfile: reader = csv.DictReader(csvfile) rows = [] for row in reader: if row['epoch'] != '': self.assertAlmostEqual(float(row['lr']), BaseCSVLoggerTest.lr) del row['lr'] rows.append(row) self.assertEqual(len(rows), len(history)) for row, hist_entry in zip(rows, history): row = {k: v for k, v in row.items() if v != ''} self.assertEqual(row.keys(), hist_entry.keys()) for k in row.keys(): if isinstance(hist_entry[k], float): self.assertAlmostEqual(float(row[k]), hist_entry[k]) else: self.assertEqual(str(row[k]), str(hist_entry[k]))
class TensorBoardGradientTrackerTest(TestCase): batch_size = 20 lr = 1e-3 num_epochs = 10 def setUp(self): torch.manual_seed(42) self.loss_function = nn.MSELoss() self.temp_dir_obj = TemporaryDirectory() # pylint: disable=not-callable self.writer = SummaryWriter(self.temp_dir_obj.name) self.writer.add_scalars = MagicMock() def tearDown(self): self.temp_dir_obj.cleanup() def test_tracking_one_layer_model(self): self.num_layer = 1 self.pytorch_network = nn.Linear(1, 1) self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=self.lr) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) keep_bias = False train_gen = some_data_generator(20) valid_gen = some_data_generator(20) tracker = TensorBoardGradientTracker(self.writer, keep_bias=keep_bias) self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[tracker]) self._test_tracking(keep_bias) def test_tracking_one_layer_model_with_bias(self): self.num_layer = 1 self.pytorch_network = nn.Linear(1, 1) self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=self.lr) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) keep_bias = True train_gen = some_data_generator(20) valid_gen = some_data_generator(20) tracker = TensorBoardGradientTracker(self.writer, keep_bias=keep_bias) self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[tracker]) self._test_tracking(keep_bias) def test_tracking_two_layers_model(self): self.num_layer = 2 self.pytorch_network = nn.Sequential(nn.Linear(1, 1), nn.Linear(1, 1)) self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=self.lr) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) keep_bias = False train_gen = some_data_generator(20) valid_gen = some_data_generator(20) tracker = TensorBoardGradientTracker(self.writer, keep_bias=keep_bias) self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[tracker]) self._test_tracking(keep_bias) def test_tracking_two_layers_shallow_model(self): self.num_layer = 2 self.pytorch_network = nn.Sequential(nn.Linear(1, 4), nn.Linear(4, 1)) self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=self.lr) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) keep_bias = False train_gen = some_data_generator(20) valid_gen = some_data_generator(20) tracker = TensorBoardGradientTracker(self.writer, keep_bias=keep_bias) self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[tracker]) self._test_tracking(keep_bias) def test_tracking_N_layers_model_with_bias(self): self.num_layer = 4 self.pytorch_network = nn.Sequential(nn.Linear(1, 1), nn.Linear(1, 1), nn.Linear(1, 1), nn.Linear(1, 1)) self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=self.lr) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) keep_bias = True train_gen = some_data_generator(20) valid_gen = some_data_generator(20) tracker = TensorBoardGradientTracker(self.writer, keep_bias=keep_bias) self.model.fit_generator(train_gen, valid_gen, epochs=self.num_epochs, steps_per_epoch=5, callbacks=[tracker]) self._test_tracking(keep_bias) def _test_tracking(self, keep_bias): expected_calls = [] for epoch in range(1, self.num_epochs + 1): layer_names = [""] if self.num_layer > 1: layer_names = [] for layer_idx in range(self.num_layer): layer_names.append("{}.".format(layer_idx)) for layer_name in layer_names: expected_calls.append( call('gradient_distributions/{}weight'.format(layer_name), {'mean': ANY}, epoch)) expected_calls.append( call('gradient_distributions/{}weight'.format(layer_name), {'mean_std_dev_up': ANY}, epoch)) expected_calls.append( call('gradient_distributions/{}weight'.format(layer_name), {'mean_std_dev_down': ANY}, epoch)) expected_calls.append( call('other_gradient_stats/{}weight'.format(layer_name), {'min': ANY}, epoch)) expected_calls.append( call('other_gradient_stats/{}weight'.format(layer_name), {'max': ANY}, epoch)) if keep_bias: expected_calls.append( call( 'gradient_distributions/{}bias'.format(layer_name), {'mean': ANY}, epoch)) expected_calls.append( call( 'gradient_distributions/{}bias'.format(layer_name), {'mean_std_dev_up': ANY}, epoch)) expected_calls.append( call( 'gradient_distributions/{}bias'.format(layer_name), {'mean_std_dev_down': ANY}, epoch)) expected_calls.append( call('other_gradient_stats/{}bias'.format(layer_name), {'min': ANY}, epoch)) expected_calls.append( call('other_gradient_stats/{}bias'.format(layer_name), {'max': ANY}, epoch)) method_calls = self.writer.add_scalars.mock_calls self.assertEqual(len(method_calls), len(expected_calls)) self.assertEqual(method_calls, expected_calls) self.assertIn(expected_calls, method_calls)
class PeriodicSaveTest(TestCase): batch_size = 20 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) self.temp_dir_obj = TemporaryDirectory() self.save_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint_{epoch}.ckpt') def tearDown(self): self.temp_dir_obj.cleanup() def test_integration(self): train_gen = some_data_generator(PeriodicSaveTest.batch_size) valid_gen = some_data_generator(PeriodicSaveTest.batch_size) saver = PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, save_best_only=True) self.model.fit_generator(train_gen, valid_gen, epochs=10, steps_per_epoch=5, callbacks=[saver]) def test_integration_with_keep_only_last_best(self): train_gen = some_data_generator(PeriodicSaveTest.batch_size) valid_gen = some_data_generator(PeriodicSaveTest.batch_size) saver = PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, save_best_only=True, keep_only_last_best=True) self.model.fit_generator(train_gen, valid_gen, epochs=10, steps_per_epoch=5, callbacks=[saver]) def test_temporary_filename_arg(self): tmp_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint.tmp.ckpt') save_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint.ckpt') train_gen = some_data_generator(PeriodicSaveTest.batch_size) valid_gen = some_data_generator(PeriodicSaveTest.batch_size) saver = PeriodicEpochSave(save_filename, monitor='val_loss', verbose=True, period=1, temporary_filename=tmp_filename) self.model.fit_generator(train_gen, valid_gen, epochs=10, steps_per_epoch=5, callbacks=[saver]) self.assertFalse(os.path.isfile(tmp_filename)) self.assertTrue(os.path.isfile(save_filename)) def test_temporary_filename_arg_with_differing_save_filename(self): epochs = 10 tmp_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint.tmp.ckpt') save_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint_{epoch}.ckpt') train_gen = some_data_generator(PeriodicSaveTest.batch_size) valid_gen = some_data_generator(PeriodicSaveTest.batch_size) saver = PeriodicEpochSave(save_filename, monitor='val_loss', verbose=True, period=1, temporary_filename=tmp_filename) self.model.fit_generator(train_gen, valid_gen, epochs=epochs, steps_per_epoch=5, callbacks=[saver]) self.assertFalse(os.path.isfile(tmp_filename)) for i in range(1, epochs + 1): self.assertTrue(os.path.isfile(save_filename.format(epoch=i))) def test_non_atomic_write(self): save_filename = os.path.join(self.temp_dir_obj.name, 'my_checkpoint.ckpt') train_gen = some_data_generator(PeriodicSaveTest.batch_size) valid_gen = some_data_generator(PeriodicSaveTest.batch_size) saver = PeriodicEpochSave(save_filename, monitor='val_loss', verbose=True, period=1, atomic_write=False) self.model.fit_generator(train_gen, valid_gen, epochs=10, steps_per_epoch=5, callbacks=[saver]) self.assertTrue(os.path.isfile(save_filename)) def test_save_best_only(self): saver = PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, save_best_only=True) val_losses = [10, 3, 8, 5, 2] has_checkpoints = [True, True, False, False, True] self._test_saver_with_val_losses(saver, val_losses, has_checkpoints) def test_save_best_only_with_keep_only_last_best(self): saver = PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, save_best_only=True, keep_only_last_best=True) val_losses = [10, 3, 8, 5, 2] has_checkpoints = [True, True, False, False, True] self._test_saver_with_val_losses(saver, val_losses, has_checkpoints, keep_only_last_best=True) def test_save_best_only_with_max(self): saver = PeriodicEpochSave(self.save_filename, monitor='val_loss', mode='max', verbose=True, save_best_only=True) val_losses = [2, 3, 8, 5, 2] has_checkpoints = [True, True, True, False, False] self._test_saver_with_val_losses(saver, val_losses, has_checkpoints) def test_save_best_only_with_max_and_keep_only_last_best(self): saver = PeriodicEpochSave( self.save_filename, monitor='val_loss', mode='max', verbose=True, save_best_only=True, keep_only_last_best=True, ) val_losses = [2, 3, 8, 5, 2] has_checkpoints = [True, True, True, False, False] self._test_saver_with_val_losses(saver, val_losses, has_checkpoints, keep_only_last_best=True) def test_periodic_with_period_of_1(self): saver = PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, period=1, save_best_only=False) val_losses = [1] * 10 has_checkpoints = [True] * 10 self._test_saver_with_val_losses(saver, val_losses, has_checkpoints) def test_periodic_with_period_of_2(self): saver = PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, period=2, save_best_only=False) val_losses = [1] * 10 has_checkpoints = [False, True] * 5 self._test_saver_with_val_losses(saver, val_losses, has_checkpoints) def test_keep_only_last_best_without_save_best_only(self): with self.assertRaises(ValueError): PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, save_best_only=False, keep_only_last_best=True) with self.assertRaises(ValueError): PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, keep_only_last_best=True) def test_save_best_only_with_restore_best(self): checkpointer = PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, save_best_only=True, restore_best=True) val_losses = [10, 3, 8, 7, 2, 5] has_checkpoints = [True, True, False, False, True, False] self._test_saver_with_val_losses(checkpointer, val_losses, has_checkpoints) self.assertEqual(5, checkpointer.restored_epoch_number) self.assertEqual(5, checkpointer.last_saved_epoch_number) def test_restore_best_without_save_best_only(self): with self.assertRaises(ValueError): PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, save_best_only=False, restore_best=True) with self.assertRaises(ValueError): PeriodicEpochSave(self.save_filename, monitor='val_loss', verbose=True, restore_best=True) def _test_saver_with_val_losses(self, saver, val_losses, has_checkpoints, keep_only_last_best=False): generator = some_data_generator(PeriodicSaveTest.batch_size) best_checkpoint_filenames = [] saver.set_params({'epochs': len(val_losses), 'steps': 1}) saver.set_model(self.model) saver.on_train_begin({}) for epoch, (val_loss, has_checkpoint) in enumerate( zip(val_losses, has_checkpoints), 1): saver.on_epoch_begin(epoch, {}) saver.on_train_batch_begin(1, {}) loss = self._update_model(generator) saver.on_train_batch_end(1, { 'batch': 1, 'size': PeriodicSaveTest.batch_size, 'loss': loss }) saver.on_epoch_end(epoch, { 'epoch': epoch, 'loss': loss, 'val_loss': val_loss }) filename = self.save_filename.format(epoch=epoch) self.assertEqual(has_checkpoint, os.path.isfile(filename)) if has_checkpoint: self.assertEqual(f'{epoch}\n', open(filename, 'r', encoding='utf-8').read()) best_checkpoint_filenames.append(os.path.realpath(filename)) files = [ os.path.realpath(os.path.join(self.temp_dir_obj.name, f)) for f in os.listdir(self.temp_dir_obj.name) ] if keep_only_last_best: self.assertEqual(1, len(files)) self.assertEqual(files[0], best_checkpoint_filenames[-1]) else: best_checkpoint_filenames = set(best_checkpoint_filenames) self.assertEqual(len(best_checkpoint_filenames), len(files)) self.assertEqual(best_checkpoint_filenames, set(files)) saver.on_train_end({}) def _update_model(self, generator): self.pytorch_network.zero_grad() x, y = next(generator) pred_y = self.pytorch_network(x) loss = self.loss_function(pred_y, y) loss.backward() self.optimizer.step() return float(loss)
class LambdaTest(ModelFittingTestCase): def setUp(self): super().setUp() torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.Adam(self.pytorch_network.parameters(), lr=1e-3) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) def test_integration_zero_args(self): lambda_callback = LambdaCallback() train_generator = some_data_tensor_generator(LambdaTest.batch_size) valid_generator = some_data_tensor_generator(LambdaTest.batch_size) test_generator = some_data_tensor_generator(LambdaTest.batch_size) self.model.fit_generator( train_generator, valid_generator, epochs=LambdaTest.epochs, steps_per_epoch=LambdaTest.steps_per_epoch, validation_steps=LambdaTest.steps_per_epoch, callbacks=[lambda_callback], ) num_steps = 10 self.model.evaluate_generator(test_generator, steps=num_steps, callbacks=[lambda_callback]) def test_with_only_on_epoch_end_arg(self): on_epoch_end = Mock() lambda_callback = LambdaCallback(on_epoch_end=on_epoch_end) train_generator = some_data_tensor_generator(LambdaTest.batch_size) valid_generator = some_data_tensor_generator(LambdaTest.batch_size) test_generator = some_data_tensor_generator(LambdaTest.batch_size) logs = self.model.fit_generator( train_generator, valid_generator, epochs=LambdaTest.epochs, steps_per_epoch=LambdaTest.steps_per_epoch, validation_steps=LambdaTest.steps_per_epoch, callbacks=[lambda_callback], ) num_steps = 10 self.model.evaluate_generator(test_generator, steps=num_steps, callbacks=[lambda_callback]) expected_calls = [ call(epoch_number, log) for epoch_number, log in enumerate(logs, 1) ] actual_calls = on_epoch_end.mock_calls self.assertEqual(len(expected_calls), len(actual_calls)) self.assertEqual(expected_calls, actual_calls) def test_lambda_test_calls(self): lambda_callback, mock_calls = self._get_lambda_callback_with_mock_args( ) num_steps = 10 generator = some_data_tensor_generator(LambdaTest.batch_size) self.model.evaluate_generator( generator, steps=num_steps, callbacks=[lambda_callback, self.mock_callback]) expected_calls = self.mock_callback.method_calls[2:] actual_calls = mock_calls.method_calls self.assertEqual(len(expected_calls), len(actual_calls)) self.assertEqual(expected_calls, actual_calls) def test_lambda_train_calls(self): lambda_callback, mock_calls = self._get_lambda_callback_with_mock_args( ) train_generator = some_data_tensor_generator(LambdaTest.batch_size) valid_generator = some_data_tensor_generator(LambdaTest.batch_size) self.model.fit_generator( train_generator, valid_generator, epochs=LambdaTest.epochs, steps_per_epoch=LambdaTest.steps_per_epoch, validation_steps=LambdaTest.steps_per_epoch, callbacks=[lambda_callback, self.mock_callback], ) expected_calls = self.mock_callback.method_calls[2:] actual_calls = mock_calls.method_calls self.assertEqual(len(expected_calls), len(actual_calls)) self.assertEqual(expected_calls, actual_calls) def _get_lambda_callback_with_mock_args(self): mock_callback = Mock(spec=Callback()) lambda_callback = LambdaCallback( on_epoch_begin=mock_callback.on_epoch_begin, on_epoch_end=mock_callback.on_epoch_end, on_train_batch_begin=mock_callback.on_train_batch_begin, on_train_batch_end=mock_callback.on_train_batch_end, on_valid_batch_begin=mock_callback.on_valid_batch_begin, on_valid_batch_end=mock_callback.on_valid_batch_end, on_test_batch_begin=mock_callback.on_test_batch_begin, on_test_batch_end=mock_callback.on_test_batch_end, on_train_begin=mock_callback.on_train_begin, on_train_end=mock_callback.on_train_end, on_valid_begin=mock_callback.on_valid_begin, on_valid_end=mock_callback.on_valid_end, on_test_begin=mock_callback.on_test_begin, on_test_end=mock_callback.on_test_end, on_backward_end=mock_callback.on_backward_end, ) return lambda_callback, mock_callback
class ModelMultiIOTest(ModelFittingTestCase): def setUp(self): super().setUp() torch.manual_seed(42) self.pytorch_network = MultiIOModel(num_input=2, num_output=2) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.model = Model( self.pytorch_network, self.optimizer, lambda y_pred, y_true: self.loss_function(y_pred[0], y_true[0]) + self.loss_function(y_pred[1], y_true[1]), batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics) def test_fitting_tensor_generator_multi_io(self): train_generator = some_data_tensor_generator_multi_io(ModelMultiIOTest.batch_size) valid_generator = some_data_tensor_generator_multi_io(ModelMultiIOTest.batch_size) logs = self.model.fit_generator(train_generator, valid_generator, epochs=ModelMultiIOTest.epochs, steps_per_epoch=ModelMultiIOTest.steps_per_epoch, validation_steps=ModelMultiIOTest.steps_per_epoch, callbacks=[self.mock_callback]) params = {'epochs': ModelMultiIOTest.epochs, 'steps': ModelMultiIOTest.steps_per_epoch} self._test_callbacks_train(params, logs) def test_fitting_with_tensor_multi_io(self): train_real_steps_per_epoch = 30 train_batch_size = ModelMultiIOTest.batch_size train_final_batch_missing_samples = 7 train_size = train_real_steps_per_epoch * train_batch_size - \ train_final_batch_missing_samples train_x = (torch.rand(train_size, 1), torch.rand(train_size, 1)) train_y = (torch.rand(train_size, 1), torch.rand(train_size, 1)) valid_real_steps_per_epoch = 10 # valid_batch_size will be the same as train_batch_size in the fit method. valid_batch_size = train_batch_size valid_final_batch_missing_samples = 3 valid_size = valid_real_steps_per_epoch * valid_batch_size - \ valid_final_batch_missing_samples valid_x = (torch.rand(valid_size, 1), torch.rand(valid_size, 1)) valid_y = (torch.rand(valid_size, 1), torch.rand(valid_size, 1)) logs = self.model.fit(train_x, train_y, validation_data=(valid_x, valid_y), epochs=ModelMultiIOTest.epochs, batch_size=train_batch_size, steps_per_epoch=None, validation_steps=None, callbacks=[self.mock_callback]) params = {'epochs': ModelMultiIOTest.epochs, 'steps': train_real_steps_per_epoch} self._test_callbacks_train(params, logs) def test_tensor_train_on_batch_multi_io(self): x1 = torch.rand(ModelMultiIOTest.batch_size, 1) x2 = torch.rand(ModelMultiIOTest.batch_size, 1) y1 = torch.rand(ModelMultiIOTest.batch_size, 1) y2 = torch.rand(ModelMultiIOTest.batch_size, 1) loss = self.model.train_on_batch((x1, x2), (y1, y2)) self.assertEqual(type(loss), float) def test_ndarray_train_on_batch_multi_io(self): x1 = np.random.rand(ModelMultiIOTest.batch_size, 1).astype(np.float32) x2 = np.random.rand(ModelMultiIOTest.batch_size, 1).astype(np.float32) y1 = np.random.rand(ModelMultiIOTest.batch_size, 1).astype(np.float32) y2 = np.random.rand(ModelMultiIOTest.batch_size, 1).astype(np.float32) loss = self.model.train_on_batch((x1, x2), (y1, y2)) self.assertEqual(type(loss), float) def test_evaluate_with_pred_multi_io(self): x = (torch.rand(ModelMultiIOTest.evaluate_dataset_len, 1), torch.rand(ModelMultiIOTest.evaluate_dataset_len, 1)) y = (torch.rand(ModelMultiIOTest.evaluate_dataset_len, 1), torch.rand(ModelMultiIOTest.evaluate_dataset_len, 1)) # We also test the unpacking. _, pred_y = self.model.evaluate(x, y, batch_size=ModelMultiIOTest.batch_size, return_pred=True) for pred in pred_y: self.assertEqual(pred.shape, (ModelMultiIOTest.evaluate_dataset_len, 1)) def test_tensor_evaluate_on_batch_multi_io(self): y = (torch.rand(ModelMultiIOTest.batch_size, 1), torch.rand(ModelMultiIOTest.batch_size, 1)) x = (torch.rand(ModelMultiIOTest.batch_size, 1), torch.rand(ModelMultiIOTest.batch_size, 1)) loss = self.model.evaluate_on_batch(x, y) self.assertEqual(type(loss), float) def test_predict_with_np_array_multi_io(self): x1 = np.random.rand(ModelMultiIOTest.evaluate_dataset_len, 1).astype(np.float32) x2 = np.random.rand(ModelMultiIOTest.evaluate_dataset_len, 1).astype(np.float32) x = (x1, x2) pred_y = self.model.predict(x, batch_size=ModelMultiIOTest.batch_size) for pred in pred_y: self.assertEqual(pred.shape, (ModelMultiIOTest.evaluate_dataset_len, 1)) def test_predict_generator_multi_io(self): num_steps = 10 generator = some_data_tensor_generator_multi_io(ModelMultiIOTest.batch_size) generator = (x for x, _ in generator) pred_y = self.model.predict_generator(generator, steps=num_steps) for pred in pred_y: self.assertEqual(pred.shape, (num_steps * ModelMultiIOTest.batch_size, 1)) def test_tensor_predict_on_batch_multi_io(self): x = (torch.rand(ModelMultiIOTest.batch_size, 1), torch.rand(ModelMultiIOTest.batch_size, 1)) pred_y = self.model.predict_on_batch(x) self._test_size_and_type_for_generator(pred_y, (ModelMultiIOTest.batch_size, 1))
class NotificationCallbackTest(TestCase): epochs = 10 steps_per_epoch = 5 batch_size = 20 lr = 0.01 def setUp(self) -> None: super().setUp() self.notification_callback_mock = MagicMock() self.notificator_mock = MagicMock() self.train_generator = some_data_tensor_generator(NotificationCallbackTest.batch_size) self.valid_generator = some_data_tensor_generator(NotificationCallbackTest.batch_size) torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=NotificationCallbackTest.lr) self.batch_metrics = [ some_batch_metric_1, ('custom_name', some_batch_metric_2), repeat_batch_metric, repeat_batch_metric, ] self.batch_metrics_names = [ 'some_batch_metric_1', 'custom_name', 'repeat_batch_metric1', 'repeat_batch_metric2', ] self.batch_metrics_values = [ some_metric_1_value, some_metric_2_value, repeat_batch_metric_value, repeat_batch_metric_value, ] self.epoch_metrics = [SomeConstantEpochMetric()] self.epoch_metrics_names = ['some_constant_epoch_metric'] self.epoch_metrics_values = [some_constant_epoch_metric_value] self.model = Model( self.pytorch_network, self.optimizer, self.loss_function, batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics, ) def test_givenANotificationCallback_whenTrainingLoop_thenSendNotification(self): notification_callback = NotificationCallback(notificator=self.notificator_mock) logs = self.model.fit_generator( self.train_generator, self.valid_generator, epochs=NotificationCallbackTest.epochs, steps_per_epoch=NotificationCallbackTest.steps_per_epoch, validation_steps=NotificationCallbackTest.steps_per_epoch, callbacks=[notification_callback], ) self._test_notificator_call(logs) def test_givenANotificationCallbackWithExperimentName_whenTrainingLoop_thenSendNotificationWithExperimentName(self): a_experiment_name = "A experiment name" notification_callback = NotificationCallback( notificator=self.notificator_mock, experiment_name=a_experiment_name ) logs = self.model.fit_generator( self.train_generator, self.valid_generator, epochs=NotificationCallbackTest.epochs, steps_per_epoch=NotificationCallbackTest.steps_per_epoch, validation_steps=NotificationCallbackTest.steps_per_epoch, callbacks=[notification_callback], ) self._test_notificator_call(logs, experiment_name=a_experiment_name) def _test_notificator_call(self, logs: Dict, experiment_name=None): experiment_name_text = f" for {experiment_name}" if experiment_name is not None else "" call_list = [] call_list.append(call.send_notification('', subject=f'Start of the training{experiment_name_text}.')) for batch_log in logs: formatted_log_data = " ".join([f"{key}: {value}\n" for key, value in batch_log.items()]) call_list.append( call.send_notification( f"Here the epoch metrics: \n{formatted_log_data}", subject=f"Epoch {batch_log['epoch']} is done{experiment_name_text}.", ) ) call_list.append(call.send_notification('', subject=f'End of the training{experiment_name_text}.')) method_calls = self.notificator_mock.method_calls self.assertEqual(len(method_calls), len(call_list)) # for set_model and set param self.assertEqual(method_calls, call_list)
class DelayCallbackTest(TestCase): epochs = 10 steps_per_epoch = 5 batch_size = 20 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) self.mock_callback = MagicMock(spec=Callback) self.delay_callback = DelayCallback(self.mock_callback) self.train_dict = {'loss': ANY, 'time': ANY} self.log_dict = {'loss': ANY, 'val_loss': ANY, 'time': ANY} def test_epoch_delay(self): epoch_delay = 4 delay_callback = DelayCallback(self.mock_callback, epoch_delay=epoch_delay) train_generator = some_data_generator(DelayCallbackTest.batch_size) valid_generator = some_data_generator(DelayCallbackTest.batch_size) self.model.fit_generator( train_generator, valid_generator, epochs=DelayCallbackTest.epochs, steps_per_epoch=DelayCallbackTest.steps_per_epoch, validation_steps=DelayCallbackTest.steps_per_epoch, callbacks=[delay_callback]) params = { 'epochs': DelayCallbackTest.epochs, 'steps': DelayCallbackTest.steps_per_epoch, 'valid_steps': DelayCallbackTest.steps_per_epoch } call_list = [] call_list.append(call.on_train_begin({})) for epoch in range(epoch_delay + 1, DelayCallbackTest.epochs + 1): call_list.append(call.on_epoch_begin(epoch, {})) for step in range(1, params['steps'] + 1): call_list.append(call.on_train_batch_begin(step, {})) call_list.append(call.on_backward_end(step)) call_list.append( call.on_train_batch_end( step, { 'batch': step, 'size': DelayCallbackTest.batch_size, **self.train_dict })) call_list.append( call.on_epoch_end(epoch, { 'epoch': epoch, **self.log_dict })) call_list.append(call.on_train_end({})) method_calls = self.mock_callback.method_calls self.assertIn(call.set_model(self.model), method_calls[:2]) self.assertIn(call.set_params(params), method_calls[:2]) self.assertEqual(len(method_calls), len(call_list) + 2) self.assertEqual(method_calls[2:], call_list) def test_batch_delay_in_middle_of_epoch(self): self._test_batch_delay(epoch_delay=5, batch_in_epoch_delay=3) def test_batch_delay_at_begin_of_epoch(self): self._test_batch_delay(epoch_delay=5, batch_in_epoch_delay=0) def test_batch_delay_when_no_delay(self): self._test_batch_delay(epoch_delay=0, batch_in_epoch_delay=0) def _test_batch_delay(self, epoch_delay, batch_in_epoch_delay): batch_delay = epoch_delay * DelayCallbackTest.steps_per_epoch + batch_in_epoch_delay delay_callback = DelayCallback(self.mock_callback, batch_delay=batch_delay) train_generator = some_data_generator(DelayCallbackTest.batch_size) valid_generator = some_data_generator(DelayCallbackTest.batch_size) self.model.fit_generator( train_generator, valid_generator, epochs=DelayCallbackTest.epochs, steps_per_epoch=DelayCallbackTest.steps_per_epoch, validation_steps=DelayCallbackTest.steps_per_epoch, callbacks=[delay_callback]) params = { 'epochs': DelayCallbackTest.epochs, 'steps': DelayCallbackTest.steps_per_epoch, 'valid_steps': DelayCallbackTest.steps_per_epoch } call_list = [] call_list.append(call.on_train_begin({})) for epoch in range(epoch_delay + 1, DelayCallbackTest.epochs + 1): call_list.append(call.on_epoch_begin(epoch, {})) start_step = batch_in_epoch_delay + 1 if epoch == epoch_delay + 1 else 1 for step in range(start_step, params['steps'] + 1): call_list.append(call.on_train_batch_begin(step, {})) call_list.append(call.on_backward_end(step)) call_list.append( call.on_train_batch_end( step, { 'batch': step, 'size': DelayCallbackTest.batch_size, **self.train_dict })) call_list.append( call.on_epoch_end(epoch, { 'epoch': epoch, **self.log_dict })) call_list.append(call.on_train_end({})) method_calls = self.mock_callback.method_calls self.assertIn(call.set_model(self.model), method_calls[:2]) self.assertIn(call.set_params(params), method_calls[:2]) self.assertEqual(len(method_calls), len(call_list) + 2) self.assertEqual(method_calls[2:], call_list)
class BestModelRestoreTest(TestCase): batch_size = 20 def setUp(self): torch.manual_seed(42) self.pytorch_network = nn.Linear(1, 1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function) def test_integration(self): train_gen = some_data_generator(20) valid_gen = some_data_generator(20) model_restore = BestModelRestore(monitor='val_loss', verbose=True) self.model.fit_generator(train_gen, valid_gen, epochs=10, steps_per_epoch=5, callbacks=[model_restore]) def test_basic_restore(self): model_restore = BestModelRestore(monitor='val_loss') val_losses = [3, 2, 8, 5, 4] best_epoch = 2 self._test_restore_with_val_losses(model_restore, val_losses, best_epoch) def test_save_best_only_with_max(self): model_restore = BestModelRestore(monitor='val_loss', mode='max') val_losses = [3, 2, 8, 5, 4] best_epoch = 3 self._test_restore_with_val_losses(model_restore, val_losses, best_epoch) def _test_restore_with_val_losses(self, checkpointer, val_losses, best_epoch): generator = some_data_generator(BestModelRestoreTest.batch_size) best_epoch_weights = None checkpointer.set_params({'epochs': len(val_losses), 'steps': 1}) checkpointer.set_model(self.model) checkpointer.on_train_begin({}) for epoch, val_loss in enumerate(val_losses, 1): checkpointer.on_epoch_begin(epoch, {}) checkpointer.on_train_batch_begin(1, {}) loss = self._update_model(generator) checkpointer.on_train_batch_end(1, {'batch': 1, 'size': BestModelRestoreTest.batch_size, 'loss': loss}) checkpointer.on_epoch_end(epoch, {'epoch': epoch, 'loss': loss, 'val_loss': val_loss}) if epoch == best_epoch: best_epoch_weights = torch_to_numpy(self.model.get_weight_copies()) checkpointer.on_train_end({}) final_weights = torch_to_numpy(self.model.get_weight_copies()) self.assertEqual(best_epoch_weights, final_weights) def _update_model(self, generator): self.pytorch_network.zero_grad() x, y = next(generator) pred_y = self.pytorch_network(x) loss = self.loss_function(pred_y, y) loss.backward() self.optimizer.step() return float(loss)
class ModelMultiDictIOTest(ModelFittingTestCase): def setUp(self): super().setUp() torch.manual_seed(42) self.pytorch_network = DictIOModel(['x1', 'x2'], ['y1', 'y2']) self.loss_function = dict_mse_loss self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.model = Model(self.pytorch_network, self.optimizer, self.loss_function, batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics) def test_fitting_tensor_generator_multi_dict_io(self): train_generator = some_data_tensor_generator_dict_io(ModelMultiDictIOTest.batch_size) valid_generator = some_data_tensor_generator_dict_io(ModelMultiDictIOTest.batch_size) logs = self.model.fit_generator(train_generator, valid_generator, epochs=ModelMultiDictIOTest.epochs, steps_per_epoch=ModelMultiDictIOTest.steps_per_epoch, validation_steps=ModelMultiDictIOTest.steps_per_epoch, callbacks=[self.mock_callback]) params = { 'epochs': ModelMultiDictIOTest.epochs, 'steps': ModelMultiDictIOTest.steps_per_epoch, 'valid_steps': ModelMultiDictIOTest.steps_per_epoch } self._test_callbacks_train(params, logs, valid_steps=ModelMultiDictIOTest.steps_per_epoch) def test_tensor_train_on_batch_multi_dict_io(self): x, y = get_batch(ModelMultiDictIOTest.batch_size) loss = self.model.train_on_batch(x, y) self.assertEqual(type(loss), float) def test_train_on_batch_with_pred_multi_dict_io(self): x, y = get_batch(ModelMultiDictIOTest.batch_size) loss, pred_y = self.model.train_on_batch(x, y, return_pred=True) self.assertEqual(type(loss), float) for value in pred_y.values(): self.assertEqual(value.shape, (ModelMultiDictIOTest.batch_size, 1)) def test_ndarray_train_on_batch_multi_dict_io(self): x1 = np.random.rand(ModelMultiDictIOTest.batch_size, 1).astype(np.float32) x2 = np.random.rand(ModelMultiDictIOTest.batch_size, 1).astype(np.float32) y1 = np.random.rand(ModelMultiDictIOTest.batch_size, 1).astype(np.float32) y2 = np.random.rand(ModelMultiDictIOTest.batch_size, 1).astype(np.float32) x, y = dict(x1=x1, x2=x2), dict(y1=y1, y2=y2) loss = self.model.train_on_batch(x, y) self.assertEqual(type(loss), float) def test_evaluate_generator_multi_dict_io(self): num_steps = 10 generator = some_data_tensor_generator_dict_io(ModelMultiDictIOTest.batch_size) loss, pred_y = self.model.evaluate_generator(generator, steps=num_steps, return_pred=True) self.assertEqual(type(loss), float) self._test_size_and_type_for_generator(pred_y, (num_steps * ModelMultiDictIOTest.batch_size, 1)) def test_tensor_evaluate_on_batch_multi_dict_io(self): x, y = get_batch(ModelMultiDictIOTest.batch_size) loss = self.model.evaluate_on_batch(x, y) self.assertEqual(type(loss), float) def test_predict_generator_multi_dict_io(self): num_steps = 10 generator = some_data_tensor_generator_dict_io(ModelMultiDictIOTest.batch_size) generator = (x for x, _ in generator) pred_y = self.model.predict_generator(generator, steps=num_steps) self._test_size_and_type_for_generator(pred_y, (num_steps * ModelMultiDictIOTest.batch_size, 1)) def test_tensor_predict_on_batch_multi_dict_io(self): x1 = torch.rand(ModelMultiDictIOTest.batch_size, 1) x2 = torch.rand(ModelMultiDictIOTest.batch_size, 1) pred_y = self.model.predict_on_batch(dict(x1=x1, x2=x2)) self._test_size_and_type_for_generator(pred_y, (ModelMultiDictIOTest.batch_size, 1))
class ModelMultiInputTest(ModelFittingTestCase): def setUp(self): super().setUp() torch.manual_seed(42) self.pytorch_network = MultiIOModel(num_input=1, num_output=1) self.loss_function = nn.MSELoss() self.optimizer = torch.optim.SGD(self.pytorch_network.parameters(), lr=1e-3) self.model = Model( self.pytorch_network, self.optimizer, self.loss_function, batch_metrics=self.batch_metrics, epoch_metrics=self.epoch_metrics, ) def test_fitting_tensor_generator_multi_input(self): train_generator = some_data_tensor_generator_multi_input( ModelMultiInputTest.batch_size) valid_generator = some_data_tensor_generator_multi_input( ModelMultiInputTest.batch_size) logs = self.model.fit_generator( train_generator, valid_generator, epochs=ModelMultiInputTest.epochs, steps_per_epoch=ModelMultiInputTest.steps_per_epoch, validation_steps=ModelMultiInputTest.steps_per_epoch, callbacks=[self.mock_callback], ) params = { 'epochs': ModelMultiInputTest.epochs, 'steps': ModelMultiInputTest.steps_per_epoch, 'valid_steps': ModelMultiInputTest.steps_per_epoch, } self._test_callbacks_train( params, logs, valid_steps=ModelMultiInputTest.steps_per_epoch) def test_fitting_with_tensor_multi_input(self): train_real_steps_per_epoch = 30 train_batch_size = ModelMultiInputTest.batch_size train_final_batch_missing_samples = 7 train_size = train_real_steps_per_epoch * train_batch_size - train_final_batch_missing_samples train_x = (torch.rand(train_size, 1), torch.rand(train_size, 1)) train_y = torch.rand(train_size, 1) valid_real_steps_per_epoch = 10 # valid_batch_size will be the same as train_batch_size in the fit method. valid_batch_size = train_batch_size valid_final_batch_missing_samples = 3 valid_size = valid_real_steps_per_epoch * valid_batch_size - valid_final_batch_missing_samples valid_x = (torch.rand(valid_size, 1), torch.rand(valid_size, 1)) valid_y = torch.rand(valid_size, 1) logs = self.model.fit( train_x, train_y, validation_data=(valid_x, valid_y), epochs=ModelMultiInputTest.epochs, batch_size=train_batch_size, steps_per_epoch=None, validation_steps=None, callbacks=[self.mock_callback], ) params = { 'epochs': ModelMultiInputTest.epochs, 'steps': train_real_steps_per_epoch, 'valid_steps': valid_real_steps_per_epoch, } self._test_callbacks_train(params, logs) def test_tensor_train_on_batch_multi_input(self): x1 = torch.rand(ModelMultiInputTest.batch_size, 1) x2 = torch.rand(ModelMultiInputTest.batch_size, 1) y = torch.rand(ModelMultiInputTest.batch_size, 1) loss = self.model.train_on_batch((x1, x2), y) self.assertEqual(type(loss), float) def test_train_on_batch_with_pred_multi_input(self): x1 = torch.rand(ModelMultiInputTest.batch_size, 1) x2 = torch.rand(ModelMultiInputTest.batch_size, 1) y = torch.rand(ModelMultiInputTest.batch_size, 1) loss, pred_y = self.model.train_on_batch((x1, x2), y, return_pred=True) self.assertEqual(type(loss), float) self.assertEqual(pred_y.shape, (ModelMultiInputTest.batch_size, 1)) def test_ndarray_train_on_batch_multi_input(self): x1 = np.random.rand(ModelMultiInputTest.batch_size, 1).astype(np.float32) x2 = np.random.rand(ModelMultiInputTest.batch_size, 1).astype(np.float32) y = np.random.rand(ModelMultiInputTest.batch_size, 1).astype(np.float32) loss = self.model.train_on_batch((x1, x2), y) self.assertEqual(type(loss), float) def test_evaluate_multi_input(self): x = ( torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1), torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1), ) y = torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1) loss = self.model.evaluate(x, y, batch_size=ModelMultiInputTest.batch_size) self.assertEqual(type(loss), float) def test_evaluate_with_pred_multi_input(self): x = ( torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1), torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1), ) y = torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1) # We also test the unpacking. _, pred_y = self.model.evaluate( x, y, batch_size=ModelMultiInputTest.batch_size, return_pred=True) self.assertEqual(pred_y.shape, (ModelMultiInputTest.evaluate_dataset_len, 1)) def test_evaluate_with_np_array_multi_input(self): x1 = np.random.rand(ModelMultiInputTest.evaluate_dataset_len, 1).astype(np.float32) x2 = np.random.rand(ModelMultiInputTest.evaluate_dataset_len, 1).astype(np.float32) x = (x1, x2) y = np.random.rand(ModelMultiInputTest.evaluate_dataset_len, 1).astype(np.float32) loss, pred_y = self.model.evaluate( x, y, batch_size=ModelMultiInputTest.batch_size, return_pred=True) self.assertEqual(type(loss), float) self.assertEqual(pred_y.shape, (ModelMultiInputTest.evaluate_dataset_len, 1)) def test_evaluate_data_loader_multi_input(self): x1 = torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1) x2 = torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1) y = torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1) dataset = TensorDataset((x1, x2), y) generator = DataLoader(dataset, ModelMultiInputTest.batch_size) loss, pred_y = self.model.evaluate_generator(generator, return_pred=True) self.assertEqual(type(loss), float) self.assertEqual(pred_y.shape, (ModelMultiInputTest.evaluate_dataset_len, 1)) def test_evaluate_generator_multi_input(self): num_steps = 10 generator = some_data_tensor_generator_multi_input( ModelMultiInputTest.batch_size) loss, pred_y = self.model.evaluate_generator(generator, steps=num_steps, return_pred=True) self.assertEqual(type(loss), float) self.assertEqual(pred_y.shape, (num_steps * ModelMultiInputTest.batch_size, 1)) def test_tensor_evaluate_on_batch_multi_input(self): x1 = torch.rand(ModelMultiInputTest.batch_size, 1) x2 = torch.rand(ModelMultiInputTest.batch_size, 1) y = torch.rand(ModelMultiInputTest.batch_size, 1) loss = self.model.evaluate_on_batch((x1, x2), y) self.assertEqual(type(loss), float) def test_predict_multi_input(self): x = ( torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1), torch.rand(ModelMultiInputTest.evaluate_dataset_len, 1), ) pred_y = self.model.predict(x, batch_size=ModelMultiInputTest.batch_size) self.assertEqual(pred_y.shape, (ModelMultiInputTest.evaluate_dataset_len, 1)) def test_predict_with_np_array_multi_input(self): x1 = np.random.rand(ModelMultiInputTest.evaluate_dataset_len, 1).astype(np.float32) x2 = np.random.rand(ModelMultiInputTest.evaluate_dataset_len, 1).astype(np.float32) x = (x1, x2) pred_y = self.model.predict(x, batch_size=ModelMultiInputTest.batch_size) self.assertEqual(pred_y.shape, (ModelMultiInputTest.evaluate_dataset_len, 1)) def test_predict_generator_multi_input(self): num_steps = 10 generator = some_data_tensor_generator_multi_input( ModelMultiInputTest.batch_size) generator = (x for x, _ in generator) pred_y = self.model.predict_generator(generator, steps=num_steps) self.assertEqual(type(pred_y), np.ndarray) self.assertEqual(pred_y.shape, (num_steps * ModelMultiInputTest.batch_size, 1)) def test_tensor_predict_on_batch_multi_input(self): x1 = torch.rand(ModelMultiInputTest.batch_size, 1) x2 = torch.rand(ModelMultiInputTest.batch_size, 1) pred_y = self.model.predict_on_batch((x1, x2)) self.assertEqual(pred_y.shape, (ModelMultiInputTest.batch_size, 1))