class LRSchedulersTest(TestCase):
batch_size = 20
epochs = 10
steps_per_epoch = 5
def setUp(self):
torch.manual_seed(42)
self.pytorch_module = nn.Linear(1, 1)
self.loss_function = nn.MSELoss()
self.optimizer = torch.optim.SGD(self.pytorch_module.parameters(),
lr=1e-3)
self.model = Model(self.pytorch_module, 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_disable_batch_size_warning(self):
import warnings
def tuple_generator(batch_size):
while True:
x1 = torch.rand(batch_size, 1)
x2 = torch.rand(batch_size, 1)
y1 = torch.rand(batch_size, 1)
y2 = torch.rand(batch_size, 1)
yield (x1, x2), (y1, y2)
class TupleModule(nn.Module):
def __init__(self):
super().__init__()
self.l1 = nn.Linear(1, 1)
self.l2 = nn.Linear(1, 1)
def forward(self, x):
# pylint: disable=arguments-differ
x1, x2 = x
return self.l1(x1), self.l2(x2)
def loss_function(y_pred, y_true):
return F.mse_loss(y_pred[0], y_true[0]) + F.mse_loss(
y_pred[1], y_true[1])
pytorch_module = TupleModule()
optimizer = torch.optim.SGD(pytorch_module.parameters(), lr=1e-3)
model = Model(pytorch_module, optimizer, loss_function)
train_generator = tuple_generator(ModelTest.batch_size)
valid_generator = tuple_generator(ModelTest.batch_size)
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
model.fit_generator(train_generator,
valid_generator,
epochs=ModelTest.epochs,
steps_per_epoch=ModelTest.steps_per_epoch,
validation_steps=ModelTest.steps_per_epoch)
num_warnings = ModelTest.steps_per_epoch * 2 * ModelTest.epochs
self.assertEqual(len(w), num_warnings)
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
warning_settings['batch_size'] = 'ignore'
model.fit_generator(train_generator,
valid_generator,
epochs=ModelTest.epochs,
steps_per_epoch=ModelTest.steps_per_epoch,
validation_steps=ModelTest.steps_per_epoch)
self.assertEqual(len(w), 0)
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)
def start_record(models, datasets, device, batch_size, representations, epochs,
metrics):
for ds in datasets:
for model_name, model in models.items():
print(f'{model_name} with {ds.name} :\n')
ds.to(device)
splitter = DataSplit(ds,
test_train_split=TEST_TRAIN_SPLIT,
val_train_split=VAL_TRAIN_SPLIT,
shuffle=True)
loaders = splitter.get_split(batch_size=batch_size)
train_loader, valid_loader, test_loader = loaders
net = model['network'](dataset=ds)
optimizer = model['optimizer'][0](net.parameters(),
**model['optimizer'][1])
base_callback = BaseCB(model_name=model_name,
dataset_name=ds.name,
records_path=RECORDS_PATH)
callbacks = [
base_callback,
DatasetCB(base_callback=base_callback,
dataset=ds,
batch_size=batch_size),
MetricsCB(base_callback=base_callback, batch_metrics=metrics),
LearningRateCB(base_callback=base_callback),
DecisionBoundariesCB(base_callback=base_callback,
dataset=ds,
representations=representations),
WeightsBiasesCB(base_callback=base_callback)
]
if model['scheduler'] is not None:
scheduler = model['scheduler'][0](**model['scheduler'][1])
callbacks.append(scheduler)
batch_metrics = [metrics_list[metric] for metric in metrics]
poutyne_model = Model(net,
optimizer,
model['loss function'],
batch_metrics=batch_metrics)
poutyne_model.to(device)
poutyne_model.fit_generator(train_loader,
valid_loader,
epochs=epochs,
callbacks=callbacks)
test_loss, test_acc = poutyne_model.evaluate_generator(test_loader)
print(f'Test:\n\tLoss: {test_loss}\n\tAccuracy: {test_acc}\n')
def main(rnn_type, n_layers, dataset, embedding, device, save_path):
train_iter, valid_iter, test_iter = dataset_factory(dataset, embedding=embedding)
embedding_dim = int(embedding.split(".")[-1][:-1])
save_path = Path(save_path) / f"{rnn_type}_{n_layers}layer_{embedding_dim}"
save_path.mkdir(parents=True, exist_ok=True)
kwargs = dict(
vocab_size=len(TEXT.vocab),
embedding_dim=embedding_dim,
hidden_dim=256,
output_dim=1,
n_layers=n_layers,
dropout=0.5,
pad_idx=TEXT.vocab.stoi[TEXT.pad_token],
rnn_type="gru",
)
with open(save_path / "kwargs.json", "w") as kwargs_file:
json.dump(kwargs, kwargs_file)
pretrained_embeddings = TEXT.vocab.vectors
network = RNN(**kwargs)
network.embedding.weight.data.copy_(pretrained_embeddings)
UNK_IDX = TEXT.vocab.stoi[TEXT.unk_token]
PAD_IDX = TEXT.vocab.stoi[TEXT.pad_token]
network.embedding.weight.data[UNK_IDX] = torch.zeros(embedding_dim)
network.embedding.weight.data[PAD_IDX] = torch.zeros(embedding_dim)
optimizer = torch.optim.Adam(network.parameters())
model = Model(
network=network,
optimizer=optimizer,
loss_function=custom_loss,
batch_metrics=[acc],
)
model.to(device)
history = model.fit_generator(
train_generator=train_iter,
valid_generator=valid_iter,
epochs=10,
callbacks=[
ModelCheckpoint(
filename=str(save_path / "model.pkl"),
save_best_only=True,
restore_best=True,
)
],
)
print(f"Model saved to {save_path}")
__import__("pudb").set_trace()
test_loss, test_acc, y_pred, y_true = model.evaluate_generator(
generator=test_iter, return_pred=True, return_ground_truth=True
)
print(f"Test Loss: {test_loss:.4f}, Test Binary Accuracy: {test_acc:.4f}")
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)
import sys
sys.path.append("..")
import torch
from torch.utils.data import DataLoader
from poutyne.framework import Model
from torch_enhance.datasets import BSDS300, Set14, Set5
from torch_enhance.models import SRCNN
from torch_enhance import metrics
scale_factor = 2
train_dataset = BSDS300(scale_factor=scale_factor, data_dir="../.data/")
val_dataset = Set14(scale_factor=scale_factor, data_dir="../.data/")
train_dataloader = DataLoader(train_dataset, batch_size=8)
val_dataloader = DataLoader(val_dataset, batch_size=2)
pytorch_network = SRCNN(scale_factor)
model = Model(pytorch_network, "sgd", "mse")
model.fit_generator(train_dataloader, val_dataloader, epochs=1)
class BestModelRestoreTest(TestCase):
batch_size = 20
def setUp(self):
torch.manual_seed(42)
self.pytorch_module = nn.Linear(1, 1)
self.loss_function = nn.MSELoss()
self.optimizer = torch.optim.SGD(self.pytorch_module.parameters(),
lr=1e-3)
self.model = Model(self.pytorch_module, 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_batch_begin(1, {})
loss = self._update_model(generator)
checkpointer.on_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_module.zero_grad()
x, y = next(generator)
pred_y = self.pytorch_module(x)
loss = self.loss_function(pred_y, y)
loss.backward()
self.optimizer.step()
return float(loss)
class Experiment:
BEST_CHECKPOINT_FILENAME = 'checkpoint_epoch_{epoch}.ckpt'
BEST_CHECKPOINT_TMP_FILENAME = 'checkpoint_epoch.tmp.ckpt'
MODEL_CHECKPOINT_FILENAME = 'checkpoint.ckpt'
MODEL_CHECKPOINT_TMP_FILENAME = 'checkpoint.tmp.ckpt'
OPTIMIZER_CHECKPOINT_FILENAME = 'checkpoint.optim'
OPTIMIZER_CHECKPOINT_TMP_FILENAME = 'checkpoint.tmp.optim'
LOG_FILENAME = 'log.tsv'
TENSORBOARD_DIRECTORY = 'tensorboard'
EPOCH_FILENAME = 'last.epoch'
EPOCH_TMP_FILENAME = 'last.tmp.epoch'
LR_SCHEDULER_FILENAME = 'lr_sched_%d.lrsched'
LR_SCHEDULER_TMP_FILENAME = 'lr_sched_%d.tmp.lrsched'
TEST_LOG_FILENAME = 'test_log.tsv'
def __init__(self,
directory,
module,
*,
device=None,
logging=True,
optimizer='sgd',
loss_function=None,
metrics=[],
monitor_metric=None,
monitor_mode=None,
type=None):
self.directory = directory
self.logging = logging
if type is not None and not type.startswith(
'classif') and not type.startswith('reg'):
raise ValueError("Invalid type '%s'" % type)
loss_function = self._get_loss_function(loss_function, module, type)
metrics = self._get_metrics(metrics, module, type)
self._set_monitor(monitor_metric, monitor_mode, type)
self.model = Model(module, optimizer, loss_function, metrics=metrics)
if device is not None:
self.model.to(device)
join_dir = lambda x: os.path.join(directory, x)
self.best_checkpoint_filename = join_dir(
Experiment.BEST_CHECKPOINT_FILENAME)
self.best_checkpoint_tmp_filename = join_dir(
Experiment.BEST_CHECKPOINT_TMP_FILENAME)
self.model_checkpoint_filename = join_dir(
Experiment.MODEL_CHECKPOINT_FILENAME)
self.model_checkpoint_tmp_filename = join_dir(
Experiment.MODEL_CHECKPOINT_TMP_FILENAME)
self.optimizer_checkpoint_filename = join_dir(
Experiment.OPTIMIZER_CHECKPOINT_FILENAME)
self.optimizer_checkpoint_tmp_filename = join_dir(
Experiment.OPTIMIZER_CHECKPOINT_TMP_FILENAME)
self.log_filename = join_dir(Experiment.LOG_FILENAME)
self.tensorboard_directory = join_dir(Experiment.TENSORBOARD_DIRECTORY)
self.epoch_filename = join_dir(Experiment.EPOCH_FILENAME)
self.epoch_tmp_filename = join_dir(Experiment.EPOCH_TMP_FILENAME)
self.lr_scheduler_filename = join_dir(Experiment.LR_SCHEDULER_FILENAME)
self.lr_scheduler_tmp_filename = join_dir(
Experiment.LR_SCHEDULER_TMP_FILENAME)
self.test_log_filename = join_dir(Experiment.TEST_LOG_FILENAME)
def _get_loss_function(self, loss_function, module, type):
if loss_function is None:
if hasattr(module, 'loss_function'):
return module.loss_function
elif type is not None:
if type.startswith('classif'):
return 'cross_entropy'
elif type.startswith('reg'):
return 'mse'
return loss_function
def _get_metrics(self, metrics, module, type):
if metrics is None or len(metrics) == 0:
if hasattr(module, 'metrics'):
return module.metrics
elif type is not None and type.startswith('classif'):
return ['accuracy']
return metrics
def _set_monitor(self, monitor_metric, monitor_mode, type):
if monitor_mode is not None and monitor_mode not in ['min', 'max']:
raise ValueError("Invalid mode '%s'" % monitor_mode)
self.monitor_metric = 'val_loss'
self.monitor_mode = 'min'
if monitor_metric is not None:
self.monitor_metric = monitor_metric
if monitor_mode is not None:
self.monitor_mode = monitor_mode
elif type is not None and type.startswith('classif'):
self.monitor_metric = 'val_acc'
self.monitor_mode = 'max'
def get_best_epoch_stats(self):
if pd is None:
warnings.warn("pandas needs to be installed to use this function.")
history = pd.read_csv(self.log_filename, sep='\t')
if self.monitor_mode == 'min':
best_epoch_index = history[self.monitor_metric].idxmin()
else:
best_epoch_index = history[self.monitor_metric].idxmax()
return history.iloc[best_epoch_index:best_epoch_index + 1]
def _warn_missing_file(self, filename):
warnings.warn("Missing checkpoint: %s." % filename)
def _load_epoch_state(self, lr_schedulers):
# pylint: disable=broad-except
initial_epoch = 1
if os.path.isfile(self.epoch_filename):
try:
with open(self.epoch_filename, 'r') as f:
initial_epoch = int(f.read()) + 1
except Exception as e:
print(e)
if os.path.isfile(self.model_checkpoint_filename):
try:
print("Loading weights from %s and starting at epoch %d." %
(self.model_checkpoint_filename, initial_epoch))
self.model.load_weights(self.model_checkpoint_filename)
except Exception as e:
print(e)
else:
self._warn_missing_file(self.model_checkpoint_filename)
if os.path.isfile(self.optimizer_checkpoint_filename):
try:
print(
"Loading optimizer state from %s and starting at epoch %d."
% (self.optimizer_checkpoint_filename, initial_epoch))
self.model.load_optimizer_state(
self.optimizer_checkpoint_filename)
except Exception as e:
print(e)
else:
self._warn_missing_file(self.optimizer_checkpoint_filename)
for i, lr_scheduler in enumerate(lr_schedulers):
filename = self.lr_scheduler_filename % i
if os.path.isfile(filename):
try:
print(
"Loading LR scheduler state from %s and starting at epoch %d."
% (filename, initial_epoch))
lr_scheduler.load_state(filename)
except Exception as e:
print(e)
else:
self._warn_missing_file(filename)
return initial_epoch
def _init_model_restoring_callbacks(self, initial_epoch, save_every_epoch):
callbacks = []
best_checkpoint = ModelCheckpoint(
self.best_checkpoint_filename,
monitor=self.monitor_metric,
mode=self.monitor_mode,
save_best_only=not save_every_epoch,
restore_best=not save_every_epoch,
verbose=not save_every_epoch,
temporary_filename=self.best_checkpoint_tmp_filename)
callbacks.append(best_checkpoint)
if save_every_epoch:
best_restore = BestModelRestore(monitor=self.monitor_metric,
mode=self.monitor_mode,
verbose=True)
callbacks.append(best_restore)
if initial_epoch > 1:
# We set the current best metric score in the ModelCheckpoint so that
# it does not save checkpoint it would not have saved if the
# optimization was not stopped.
best_epoch_stats = self.get_best_epoch_stats()
best_epoch = best_epoch_stats['epoch'].item()
best_filename = self.best_checkpoint_filename.format(
epoch=best_epoch)
if not save_every_epoch:
best_checkpoint.best_filename = best_filename
best_checkpoint.current_best = best_epoch_stats[
self.monitor_metric].item()
else:
best_restore.best_weights = torch.load(best_filename,
map_location='cpu')
best_restore.current_best = best_epoch_stats[
self.monitor_metric].item()
return callbacks
def _init_tensorboard_callbacks(self, disable_tensorboard):
tensorboard_writer = None
callbacks = []
if not disable_tensorboard:
if SummaryWriter is None:
warnings.warn(
"tensorboardX does not seem to be installed. "
"To remove this warning, set the 'disable_tensorboard' "
"flag to True.")
else:
tensorboard_writer = SummaryWriter(self.tensorboard_directory)
callbacks += [TensorBoardLogger(tensorboard_writer)]
return tensorboard_writer, callbacks
def _init_lr_scheduler_callbacks(self, lr_schedulers):
callbacks = []
if self.logging:
for i, lr_scheduler in enumerate(lr_schedulers):
filename = self.lr_scheduler_filename % i
tmp_filename = self.lr_scheduler_tmp_filename % i
callbacks += [
LRSchedulerCheckpoint(lr_scheduler,
filename,
verbose=False,
temporary_filename=tmp_filename)
]
else:
callbacks += lr_schedulers
callbacks += [
BestModelRestore(monitor=self.monitor_metric,
mode=self.monitor_mode,
verbose=True)
]
return callbacks
def train(self,
train_loader,
valid_loader=None,
*,
callbacks=[],
lr_schedulers=[],
save_every_epoch=False,
disable_tensorboard=False,
epochs=1000,
steps_per_epoch=None,
validation_steps=None,
seed=42):
if seed is not None:
# Make training deterministic.
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
# Copy callback list.
callbacks = list(callbacks)
tensorboard_writer = None
initial_epoch = 1
if self.logging:
if not os.path.exists(self.directory):
os.makedirs(self.directory)
# Restarting optimization if needed.
initial_epoch = self._load_epoch_state(lr_schedulers)
callbacks += [
CSVLogger(self.log_filename,
separator='\t',
append=initial_epoch != 1)
]
callbacks += self._init_model_restoring_callbacks(
initial_epoch, save_every_epoch)
callbacks += [
ModelCheckpoint(
self.model_checkpoint_filename,
verbose=False,
temporary_filename=self.model_checkpoint_tmp_filename)
]
callbacks += [
OptimizerCheckpoint(
self.optimizer_checkpoint_filename,
verbose=False,
temporary_filename=self.optimizer_checkpoint_tmp_filename)
]
# We save the last epoch number after the end of the epoch so that the
# _load_epoch_state() knows which epoch to restart the optimization.
callbacks += [
PeriodicSaveLambda(
lambda fd, epoch, logs: print(epoch, file=fd),
self.epoch_filename,
temporary_filename=self.epoch_tmp_filename,
open_mode='w')
]
tensorboard_writer, cb_list = self._init_tensorboard_callbacks(
disable_tensorboard)
callbacks += cb_list
# This method returns callbacks that checkpoints the LR scheduler if logging is enabled.
# Otherwise, it just returns the list of LR schedulers with a BestModelRestore callback.
callbacks += self._init_lr_scheduler_callbacks(lr_schedulers)
try:
return self.model.fit_generator(train_loader,
valid_loader,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
validation_steps=validation_steps,
initial_epoch=initial_epoch,
callbacks=callbacks)
finally:
if tensorboard_writer is not None:
tensorboard_writer.close()
def load_best_checkpoint(self, *, verbose=False):
best_epoch_stats = self.get_best_epoch_stats()
best_epoch = best_epoch_stats['epoch'].item()
if verbose:
metrics_str = ', '.join(
'%s: %g' % (metric_name, best_epoch_stats[metric_name].item())
for metric_name in best_epoch_stats.columns[2:])
print("Found best checkpoint at epoch: {}".format(best_epoch))
print(metrics_str)
self.load_checkpoint(best_epoch)
return best_epoch_stats
def load_checkpoint(self, epoch):
ckpt_filename = self.best_checkpoint_filename.format(epoch=epoch)
self.model.load_weights(ckpt_filename)
def load_last_checkpoint(self):
self.model.load_weights(self.model_checkpoint_filename)
def test(self,
test_loader,
*,
steps=None,
load_best_checkpoint=True,
load_last_checkpoint=False,
seed=42):
if seed is not None:
# Make training deterministic.
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
best_epoch_stats = None
if load_best_checkpoint:
best_epoch_stats = self.load_best_checkpoint(verbose=True)
elif load_last_checkpoint:
best_epoch_stats = self.load_last_checkpoint()
test_loss, test_metrics = self.model.evaluate_generator(test_loader,
steps=steps)
if not isinstance(test_metrics, np.ndarray):
test_metrics = np.array([test_metrics])
test_metrics_names = ['test_loss'] + \
['test_' + metric_name for metric_name in self.model.metrics_names]
test_metrics_values = np.concatenate(([test_loss], test_metrics))
test_metrics_dict = {
col: val
for col, val in zip(test_metrics_names, test_metrics_values)
}
test_metrics_str = ', '.join('%s: %g' % (col, val)
for col, val in test_metrics_dict.items())
print("On best model: %s" % test_metrics_str)
if self.logging:
test_stats = pd.DataFrame([test_metrics_values],
columns=test_metrics_names)
if best_epoch_stats is not None:
best_epoch_stats = best_epoch_stats.reset_index(drop=True)
test_stats = best_epoch_stats.join(test_stats)
test_stats.to_csv(self.test_log_filename, sep='\t', index=False)
return test_metrics_dict
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()
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
}
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_batch_begin(step, {}))
call_list.append(call.on_backward_end(step))
call_list.append(
call.on_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
}
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_batch_begin(step, {}))
call_list.append(call.on_backward_end(step))
call_list.append(
call.on_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 CSVLoggerTest(TestCase):
batch_size = 20
lr = 1e-3
def setUp(self):
torch.manual_seed(42)
self.pytorch_module = nn.Linear(1, 1)
self.loss_function = nn.MSELoss()
self.optimizer = torch.optim.SGD(self.pytorch_module.parameters(),
lr=CSVLoggerTest.lr)
self.model = Model(self.pytorch_module, 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 = CSVLogger(self.csv_filename)
history = self.model.fit_generator(train_gen,
valid_gen,
epochs=10,
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 = CSVLogger(self.csv_filename, batch_granularity=True)
history = History()
self.model.fit_generator(train_gen,
valid_gen,
epochs=10,
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 = CSVLogger(self.csv_filename)
history = self.model.fit_generator(train_gen,
valid_gen,
epochs=10,
steps_per_epoch=5,
callbacks=[logger])
logger = CSVLogger(self.csv_filename, append=True)
history2 = self.model.fit_generator(train_gen,
valid_gen,
epochs=20,
steps_per_epoch=5,
initial_epoch=10,
callbacks=[logger])
self._test_logging(history + history2)
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']), CSVLoggerTest.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]))
def run(self, constant_input, hidden_state):
constant_input.requires_grad = False
hidden_state.requires_grad = True
class DataIterator(Dataset):
"""
Required dataset class to circumvent some of poutyne's limitation.
The short version is that calling `.fit()` creates a TensorDataset
(poutyne's own version) and it checks that the first dimension of all
inputs are of same dimensions. The nature of RNNCell makes it such
that the input and the hidden state cannot be aligned on the first
dimension.
So we create our own iterator and use `.fit_generator()` instead
"""
def __init__(self, constant_input, hidden_state, batch_size=32):
self.constant_input = constant_input
self.hidden_state = hidden_state
self._batch_size = batch_size
assert self.constant_input.shape[0] == self.hidden_state.shape[1]
def __len__(self):
num_items = self.constant_input.shape[0]
l = num_items // self._batch_size
l += 1 if num_items % self._batch_size != 0 else 0
return l
def __iter__(self):
last_idx = self.constant_input.shape[0]
last_idx += last_idx % self._batch_size
for start in range(0, last_idx, self._batch_size):
end = start + self._batch_size
x = self.constant_input[start:end]
y = self.hidden_state[:, start:end]
yield (x, y), y
model = Model(
network=self._rnn_cell,
loss_function=speed_loss,
optimizer=torch.optim.Adam(params=[hidden_state], lr=self._lr),
)
model.fit_generator(
DataIterator(constant_input, hidden_state, batch_size=self._batch_size),
epochs=self._n_iter,
verbose=False,
callbacks=[
StepLR(step_size=1000, gamma=0.5),
EarlyStopping(monitor="loss", min_delta=1e-6, patience=1000),
ModelCheckpoint(
filename=NamedTemporaryFile().name,
monitor="loss",
save_best_only=True,
restore_best=True,
),
],
)
trained = hidden_state.clone().detach()
_, output = model.evaluate_generator(
DataIterator(constant_input, trained, batch_size=self._batch_size),
return_pred=True,
)
output = np.concatenate([o[0] for o in output])
if trained.device.type == "cuda":
trained = trained.detach().cpu().numpy()
hidden_state = hidden_state.detach().cpu().numpy()
else:
trained = trained.detach().numpy()
hidden_state = hidden_state.detach().numpy()
return (
hidden_state.squeeze(),
_speed_loss(np.squeeze(trained), np.squeeze(output)),
)
class ModelTest(TestCase):
# pylint: disable=too-many-public-methods
epochs = 10
steps_per_epoch = 5
batch_size = 20
evaluate_dataset_len = 107
cuda_device = int(os.environ.get('CUDA_DEVICE', 0))
def setUp(self):
torch.manual_seed(42)
self.pytorch_module = nn.Linear(1, 1)
self.loss_function = nn.MSELoss()
self.optimizer = torch.optim.SGD(self.pytorch_module.parameters(),
lr=1e-3)
self.metrics = [some_metric_1, some_metric_2]
self.metrics_names = ['some_metric_1', 'some_metric_2']
self.metrics_values = [some_metric_1_value, some_metric_2_value]
self.model = Model(self.pytorch_module,
self.optimizer,
self.loss_function,
metrics=self.metrics)
self.mock_callback = MagicMock()
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
}
self._test_fitting(params, logs)
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_fitting(params, logs, has_valid=False)
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
}
self._test_fitting(params, logs)
def test_fitting_with_data_loader(self):
# pylint: disable=too-many-locals
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
}
self._test_fitting(params, logs)
def test_fitting_with_tensor(self):
# pylint: disable=too-many-locals
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_x=valid_x,
validation_y=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
}
self._test_fitting(params, logs)
def test_fitting_with_np_array(self):
# pylint: disable=too-many-locals
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_x=valid_x,
validation_y=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
}
self._test_fitting(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_generator = SomeDataGeneratorWithLen(batch_size=15,
length=10,
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
}
self._test_fitting(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_fitting(params, logs, steps=train_real_steps_per_epoch)
def _test_fitting(self, params, logs, has_valid=True, steps=None):
if steps is None:
steps = params['steps']
self.assertEqual(len(logs), params['epochs'])
train_dict = dict(zip(self.metrics_names, self.metrics_values),
loss=ANY,
time=ANY)
if has_valid:
val_metrics_names = [
'val_' + metric_name for metric_name in self.metrics_names
]
val_dict = dict(zip(val_metrics_names, self.metrics_values),
val_loss=ANY)
log_dict = {**train_dict, **val_dict}
else:
log_dict = train_dict
for epoch, log in enumerate(logs, 1):
self.assertEqual(log, dict(log_dict, epoch=epoch))
call_list = []
call_list.append(call.on_train_begin({}))
for epoch in range(1, params['epochs'] + 1):
call_list.append(call.on_epoch_begin(epoch, {}))
for step in range(1, steps + 1):
call_list.append(call.on_batch_begin(step, {}))
call_list.append(call.on_backward_end(step))
call_list.append(
call.on_batch_end(step, {
'batch': step,
'size': ANY,
**train_dict
}))
call_list.append(
call.on_epoch_end(epoch, {
'epoch': epoch,
**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_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(),
[some_metric_1_value, some_metric_2_value])
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(),
[some_metric_1_value, some_metric_2_value])
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(),
[some_metric_1_value, some_metric_2_value])
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(),
[some_metric_1_value, some_metric_2_value])
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.
# pylint: disable=unused-variable
loss, metrics, 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_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(),
[some_metric_1_value, some_metric_2_value])
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(),
[some_metric_1_value, some_metric_2_value])
self._test_predictions_for_evaluate_and_predict_generator(pred_y)
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(),
[some_metric_1_value, some_metric_2_value])
for pred in pred_y:
self.assertEqual(type(pred), np.ndarray)
self.assertEqual(pred.shape, (ModelTest.batch_size, 1))
self.assertEqual(
np.concatenate(pred_y).shape,
(num_steps * ModelTest.batch_size, 1))
def test_evaluate_with_only_one_metric(self):
self.model = Model(self.pytorch_module,
self.optimizer,
self.loss_function,
metrics=self.metrics[:1])
x = torch.rand(ModelTest.evaluate_dataset_len, 1)
y = torch.rand(ModelTest.evaluate_dataset_len, 1)
loss, first_metric = self.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
self.model = Model(self.pytorch_module,
self.optimizer,
self.loss_function,
metrics=[F.mse_loss])
train_generator = some_data_tensor_generator(ModelTest.batch_size)
valid_generator = some_data_tensor_generator(ModelTest.batch_size)
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])
generator = some_data_tensor_generator(ModelTest.batch_size)
loss, mse = self.model.evaluate_generator(generator, steps=num_steps)
self.assertEqual(type(loss), float)
self.assertEqual(type(mse), float)
def test_evaluate_with_no_metric(self):
self.model = Model(self.pytorch_module, self.optimizer,
self.loss_function)
x = torch.rand(ModelTest.evaluate_dataset_len, 1)
y = torch.rand(ModelTest.evaluate_dataset_len, 1)
loss = self.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(),
[some_metric_1_value, some_metric_2_value])
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(),
[some_metric_1_value, some_metric_2_value])
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(),
[some_metric_1_value, some_metric_2_value])
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(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._test_predictions_for_evaluate_and_predict_generator(pred_y)
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)
for pred in pred_y:
self.assertEqual(type(pred), np.ndarray)
self.assertEqual(pred.shape, (ModelTest.batch_size, 1))
self.assertEqual(
np.concatenate(pred_y).shape,
(num_steps * ModelTest.batch_size, 1))
def _test_predictions_for_evaluate_and_predict_generator(self, pred_y):
self.assertEqual(type(pred_y), list)
remaning_example = ModelTest.evaluate_dataset_len
cur_batch_size = ModelTest.batch_size
for pred in pred_y:
self.assertEqual(type(pred), np.ndarray)
if remaning_example < ModelTest.batch_size:
cur_batch_size = remaning_example
remaning_example = 0
else:
remaning_example -= ModelTest.batch_size
self.assertEqual(pred.shape, (cur_batch_size, 1))
self.assertEqual(
np.concatenate(pred_y).shape, (ModelTest.evaluate_dataset_len, 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)
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_device(self, device):
for p in self.pytorch_module.parameters():
self.assertEqual(p.device, device)
def test_disable_batch_size_warning(self):
import warnings
def tuple_generator(batch_size):
while True:
x1 = torch.rand(batch_size, 1)
x2 = torch.rand(batch_size, 1)
y1 = torch.rand(batch_size, 1)
y2 = torch.rand(batch_size, 1)
yield (x1, x2), (y1, y2)
class TupleModule(nn.Module):
def __init__(self):
super().__init__()
self.l1 = nn.Linear(1, 1)
self.l2 = nn.Linear(1, 1)
def forward(self, x):
# pylint: disable=arguments-differ
x1, x2 = x
return self.l1(x1), self.l2(x2)
def loss_function(y_pred, y_true):
return F.mse_loss(y_pred[0], y_true[0]) + F.mse_loss(
y_pred[1], y_true[1])
pytorch_module = TupleModule()
optimizer = torch.optim.SGD(pytorch_module.parameters(), lr=1e-3)
model = Model(pytorch_module, optimizer, loss_function)
train_generator = tuple_generator(ModelTest.batch_size)
valid_generator = tuple_generator(ModelTest.batch_size)
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
model.fit_generator(train_generator,
valid_generator,
epochs=ModelTest.epochs,
steps_per_epoch=ModelTest.steps_per_epoch,
validation_steps=ModelTest.steps_per_epoch)
num_warnings = ModelTest.steps_per_epoch * 2 * ModelTest.epochs
self.assertEqual(len(w), num_warnings)
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
warning_settings['batch_size'] = 'ignore'
model.fit_generator(train_generator,
valid_generator,
epochs=ModelTest.epochs,
steps_per_epoch=ModelTest.steps_per_epoch,
validation_steps=ModelTest.steps_per_epoch)
self.assertEqual(len(w), 0)
experiment.log_parameters(params)
# create our special resnet18
cnn = resnet18(n_classes=4).to(device)
# print the model summary to show useful information
logging.info(summary(cnn, (3, 224, 244)))
# define custom optimizer and instantiace the trainer `Model`
optimizer = optim.Adam(cnn.parameters(), lr=params['lr'])
model = Model(cnn, optimizer, "cross_entropy",
batch_metrics=["accuracy"]).to(device)
# usually you want to reduce the lr on plateau and store the best model
callbacks = [
ReduceLROnPlateau(monitor="val_acc", patience=5, verbose=True),
ModelCheckpoint(str(project.checkpoint_dir /
f"{time.time()}-model.pt"),
save_best_only="True",
verbose=True),
EarlyStopping(monitor="val_acc", patience=10, mode='max'),
CometCallback(experiment)
]
model.fit_generator(
train_dl,
val_dl,
epochs=params['epochs'],
callbacks=callbacks,
)
# # get the results on the test set
# loss, test_acc = model.evaluate_generator(test_dl)
# logging.info(f'test_acc=({test_acc})')
# experiment.log_metric('test_acc', test_acc)
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 Experiment:
"""
The Experiment class provides a straightforward experimentation tool for efficient and entirely
customizable finetuning of the whole neural network training procedure with PyTorch. The
``Experiment`` object takes care of the training and testing processes while also managing to
keep traces of all pertinent information via the automatic logging option.
Args:
directory (str): Path to the experiment's working directory. Will be used for the automatic logging.
network (torch.nn.Module): A PyTorch network.
device (torch.torch.device): The device to which the model is sent. If None, the model will be
kept on its current device.
(Default value = None)
logging (bool): Whether or not to log the experiment's progress. If true, various logging
callbacks will be inserted to output training and testing stats as well as to automatically
save model checkpoints, for example. See :func:`~Experiment.train()` and :func:`~Experiment.test()`
for more details.
(Default value = True)
optimizer (Union[torch.optim.Optimizer, str]): If Pytorch Optimizer, must already be initialized.
If str, should be the optimizer's name in Pytorch (i.e. 'Adam' for torch.optim.Adam).
(Default value = 'sgd')
loss_function(Union[Callable, str]) It can be any PyTorch
loss layer or custom loss function. It can also be a string with the same name as a PyTorch
loss function (either the functional or object name). The loss function must have the signature
``loss_function(input, target)`` where ``input`` is the prediction of the network and ``target``
is the ground truth. If ``None``, will default to, in priority order, either the model's own
loss function or the default loss function associated with the ``task``.
(Default value = None)
batch_metrics (list): List of functions with the same signature as the loss function. Each metric
can be any PyTorch loss function. It can also be a string with the same name as a PyTorch
loss function (either the functional or object name). 'accuracy' (or just 'acc') is also a
valid metric. Each metric function is called on each batch of the optimization and on the
validation batches at the end of the epoch.
(Default value = None)
epoch_metrics (list): List of functions with the same signature as
:class:`~poutyne.framework.metrics.epoch_metrics.EpochMetric`
(Default value = None)
monitor_metric (str): Which metric to consider for best model performance calculation. Should be in
the format '{metric_name}' or 'val_{metric_name}' (i.e. 'val_loss'). If None, will follow the value
suggested by ``task`` or default to 'val_loss'.
.. warning:: If you do not plan on using a validation set, you must set the monitor metric to another
value.
monitor_mode (str): Which mode, either 'min' or 'max', should be used when considering the ``monitor_metric``
value. If None, will follow the value suggested by ``task`` or default to 'min'.
task (str): Any str beginning with either 'classif' or 'reg'. Specifying a ``task`` can assign default
values to the ``loss_function``, ``batch_metrics``, ``monitor_mode`` and ``monitor_mode``. For ``task``
that begins with 'reg', the only default value is the loss function that is the mean squared error. When
beginning with 'classif', the default loss function is the cross-entropy loss, the default batch metrics
will be the accuracy, the default epoch metrics will be the F1 score and the default monitoring will be
set on 'val_acc' with a 'max' mode.
(Default value = None)
Example:
Using a PyTorch DataLoader, on classification task with SGD optimizer::
import torch
from torch.utils.data import DataLoader, TensorDataset
from poutyne.framework import Experiment
num_features = 20
num_classes = 5
# Our training dataset with 800 samples.
num_train_samples = 800
train_x = torch.rand(num_train_samples, num_features)
train_y = torch.randint(num_classes, (num_train_samples, ), dtype=torch.long)
train_dataset = TensorDataset(train_x, train_y)
train_generator = DataLoader(train_dataset, batch_size=32)
# Our validation dataset with 200 samples.
num_valid_samples = 200
valid_x = torch.rand(num_valid_samples, num_features)
valid_y = torch.randint(num_classes, (num_valid_samples, ), dtype=torch.long)
valid_dataset = TensorDataset(valid_x, valid_y)
valid_generator = DataLoader(valid_dataset, batch_size=32)
# Our network
pytorch_network = torch.nn.Linear(num_features, num_train_samples)
# Intialization of our experimentation and network training
exp = Experiment('./simple_example',
pytorch_network,
optimizer='sgd',
task='classif')
exp.train(train_generator, valid_generator, epochs=5)
The above code will yield an output similar to the below lines. Note the automatic checkpoint saving
in the experiment directory when the monitored metric improved.
.. code-block:: none
Epoch 1/5 0.09s Step 25/25: loss: 6.351375, acc: 1.375000, val_loss: 6.236106, val_acc: 5.000000
Epoch 1: val_acc improved from -inf to 5.00000, saving file to ./simple_example/checkpoint_epoch_1.ckpt
Epoch 2/5 0.10s Step 25/25: loss: 6.054254, acc: 14.000000, val_loss: 5.944495, val_acc: 19.500000
Epoch 2: val_acc improved from 5.00000 to 19.50000, saving file to ./simple_example/checkpoint_epoch_2.ckpt
Epoch 3/5 0.09s Step 25/25: loss: 5.759377, acc: 22.875000, val_loss: 5.655412, val_acc: 21.000000
Epoch 3: val_acc improved from 19.50000 to 21.00000, saving file to ./simple_example/checkpoint_epoch_3.ckpt
...
Training can now easily be resumed from the best checkpoint::
exp.train(train_generator, valid_generator, epochs=10)
.. code-block:: none
Restoring model from ./simple_example/checkpoint_epoch_3.ckpt
Loading weights from ./simple_example/checkpoint.ckpt and starting at epoch 6.
Loading optimizer state from ./simple_example/checkpoint.optim and starting at epoch 6.
Epoch 6/10 0.16s Step 25/25: loss: 4.897135, acc: 22.875000, val_loss: 4.813141, val_acc: 20.500000
Epoch 7/10 0.10s Step 25/25: loss: 4.621514, acc: 22.625000, val_loss: 4.545359, val_acc: 20.500000
Epoch 8/10 0.24s Step 25/25: loss: 4.354721, acc: 23.625000, val_loss: 4.287117, val_acc: 20.500000
...
Testing is also very intuitive::
exp.test(test_generator)
.. code-block:: none
Restoring model from ./simple_example/checkpoint_epoch_9.ckpt
Found best checkpoint at epoch: 9
lr: 0.01, loss: 4.09892, acc: 23.625, val_loss: 4.04057, val_acc: 21.5
On best model: test_loss: 4.06664, test_acc: 17.5
Finally, all the pertinent metrics specified to the Experiment at each epoch are stored in a specific logging
file, found here at './simple_example/log.tsv'.
.. code-block:: none
epoch time lr loss acc val_loss val_acc
1 0.0721172170015052 0.01 6.351375141143799 1.375 6.23610631942749 5.0
2 0.0298177790245972 0.01 6.054253826141357 14.000 5.94449516296386 19.5
3 0.0637106419890187 0.01 5.759376544952392 22.875 5.65541223526001 21.0
...
"""
BEST_CHECKPOINT_FILENAME = 'checkpoint_epoch_{epoch}.ckpt'
BEST_CHECKPOINT_TMP_FILENAME = 'checkpoint_epoch.tmp.ckpt'
MODEL_CHECKPOINT_FILENAME = 'checkpoint.ckpt'
MODEL_CHECKPOINT_TMP_FILENAME = 'checkpoint.tmp.ckpt'
OPTIMIZER_CHECKPOINT_FILENAME = 'checkpoint.optim'
OPTIMIZER_CHECKPOINT_TMP_FILENAME = 'checkpoint.tmp.optim'
LOG_FILENAME = 'log.tsv'
LOG_TMP_FILENAME = 'log.tmp.tsv'
TENSORBOARD_DIRECTORY = 'tensorboard'
EPOCH_FILENAME = 'last.epoch'
EPOCH_TMP_FILENAME = 'last.tmp.epoch'
LR_SCHEDULER_FILENAME = 'lr_sched_%d.lrsched'
LR_SCHEDULER_TMP_FILENAME = 'lr_sched_%d.tmp.lrsched'
TEST_LOG_FILENAME = 'test_log.tsv'
def __init__(self,
directory,
network,
*,
device=None,
logging=True,
optimizer='sgd',
loss_function=None,
batch_metrics=None,
epoch_metrics=None,
monitor_metric=None,
monitor_mode=None,
task=None):
self.directory = directory
self.logging = logging
if task is not None and not task.startswith(
'classif') and not task.startswith('reg'):
raise ValueError("Invalid task '%s'" % task)
batch_metrics = [] if batch_metrics is None else batch_metrics
epoch_metrics = [] if epoch_metrics is None else epoch_metrics
loss_function = self._get_loss_function(loss_function, network, task)
batch_metrics = self._get_batch_metrics(batch_metrics, network, task)
epoch_metrics = self._get_epoch_metrics(epoch_metrics, network, task)
self._set_monitor(monitor_metric, monitor_mode, task)
self.model = Model(network,
optimizer,
loss_function,
batch_metrics=batch_metrics,
epoch_metrics=epoch_metrics)
if device is not None:
self.model.to(device)
self.best_checkpoint_filename = self.get_path(
Experiment.BEST_CHECKPOINT_FILENAME)
self.best_checkpoint_tmp_filename = self.get_path(
Experiment.BEST_CHECKPOINT_TMP_FILENAME)
self.model_checkpoint_filename = self.get_path(
Experiment.MODEL_CHECKPOINT_FILENAME)
self.model_checkpoint_tmp_filename = self.get_path(
Experiment.MODEL_CHECKPOINT_TMP_FILENAME)
self.optimizer_checkpoint_filename = self.get_path(
Experiment.OPTIMIZER_CHECKPOINT_FILENAME)
self.optimizer_checkpoint_tmp_filename = self.get_path(
Experiment.OPTIMIZER_CHECKPOINT_TMP_FILENAME)
self.log_filename = self.get_path(Experiment.LOG_FILENAME)
self.log_tmp_filename = self.get_path(Experiment.LOG_TMP_FILENAME)
self.tensorboard_directory = self.get_path(
Experiment.TENSORBOARD_DIRECTORY)
self.epoch_filename = self.get_path(Experiment.EPOCH_FILENAME)
self.epoch_tmp_filename = self.get_path(Experiment.EPOCH_TMP_FILENAME)
self.lr_scheduler_filename = self.get_path(
Experiment.LR_SCHEDULER_FILENAME)
self.lr_scheduler_tmp_filename = self.get_path(
Experiment.LR_SCHEDULER_TMP_FILENAME)
self.test_log_filename = self.get_path(Experiment.TEST_LOG_FILENAME)
def get_path(self, *paths):
"""
Returns the path inside the experiment directory.
"""
return os.path.join(self.directory, *paths)
def _get_loss_function(self, loss_function, network, task):
if loss_function is None:
if hasattr(network, 'loss_function'):
return network.loss_function
if task is not None:
if task.startswith('classif'):
return 'cross_entropy'
if task.startswith('reg'):
return 'mse'
return loss_function
def _get_batch_metrics(self, batch_metrics, network, task):
if batch_metrics is None or len(batch_metrics) == 0:
if hasattr(network, 'batch_metrics'):
return network.batch_metrics
if task is not None and task.startswith('classif'):
return ['accuracy']
return batch_metrics
def _get_epoch_metrics(self, epoch_metrics, network, task):
if epoch_metrics is None or len(epoch_metrics) == 0:
if hasattr(network, 'epoch_metrics'):
return network.epoch_metrics
if task is not None and task.startswith('classif'):
return ['f1']
return epoch_metrics
def _set_monitor(self, monitor_metric, monitor_mode, task):
if monitor_mode is not None and monitor_mode not in ['min', 'max']:
raise ValueError("Invalid mode '%s'" % monitor_mode)
self.monitor_metric = 'val_loss'
self.monitor_mode = 'min'
if monitor_metric is not None:
self.monitor_metric = monitor_metric
if monitor_mode is not None:
self.monitor_mode = monitor_mode
elif task is not None and task.startswith('classif'):
self.monitor_metric = 'val_acc'
self.monitor_mode = 'max'
def get_best_epoch_stats(self):
"""
Returns all computed statistics corresponding to the best epoch according to the
``monitor_metric`` and ``monitor_mode`` attributes.
Returns:
dict where each key is a column name in the logging output file
and values are the ones found at the best epoch.
"""
if pd is None:
raise ImportError(
"pandas needs to be installed to use this function.")
history = pd.read_csv(self.log_filename, sep='\t')
if self.monitor_mode == 'min':
best_epoch_index = history[self.monitor_metric].idxmin()
else:
best_epoch_index = history[self.monitor_metric].idxmax()
return history.iloc[best_epoch_index:best_epoch_index + 1]
def get_saved_epochs(self):
"""
Returns a pandas DataFrame which each row corresponds to an epoch having
a saved checkpoint.
Returns:
pandas DataFrame which each row corresponds to an epoch having a saved
checkpoint.
"""
if pd is None:
raise ImportError(
"pandas needs to be installed to use this function.")
history = pd.read_csv(self.log_filename, sep='\t')
metrics = history[self.monitor_metric].tolist()
if self.monitor_mode == 'min':
monitor_op = lambda x, y: x < y
current_best = float('Inf')
elif self.monitor_mode == 'max':
monitor_op = lambda x, y: x > y
current_best = -float('Inf')
saved_epoch_indices = []
for i, metric in enumerate(metrics):
if monitor_op(metric, current_best):
current_best = metric
saved_epoch_indices.append(i)
return history.iloc[saved_epoch_indices]
def _warn_missing_file(self, filename):
warnings.warn("Missing checkpoint: %s." % filename)
def _load_epoch_state(self, lr_schedulers):
# pylint: disable=broad-except
initial_epoch = 1
if os.path.isfile(self.epoch_filename):
try:
with open(self.epoch_filename, 'r') as f:
initial_epoch = int(f.read()) + 1
except Exception as e:
print(e)
if os.path.isfile(self.model_checkpoint_filename):
try:
print("Loading weights from %s and starting at epoch %d." %
(self.model_checkpoint_filename, initial_epoch))
self.model.load_weights(self.model_checkpoint_filename)
except Exception as e:
print(e)
else:
self._warn_missing_file(self.model_checkpoint_filename)
if os.path.isfile(self.optimizer_checkpoint_filename):
try:
print(
"Loading optimizer state from %s and starting at epoch %d."
% (self.optimizer_checkpoint_filename, initial_epoch))
self.model.load_optimizer_state(
self.optimizer_checkpoint_filename)
except Exception as e:
print(e)
else:
self._warn_missing_file(self.optimizer_checkpoint_filename)
for i, lr_scheduler in enumerate(lr_schedulers):
filename = self.lr_scheduler_filename % i
if os.path.isfile(filename):
try:
print(
"Loading LR scheduler state from %s and starting at epoch %d."
% (filename, initial_epoch))
lr_scheduler.load_state(filename)
except Exception as e:
print(e)
else:
self._warn_missing_file(filename)
return initial_epoch
def _init_model_restoring_callbacks(self, initial_epoch, save_every_epoch):
callbacks = []
best_checkpoint = ModelCheckpoint(
self.best_checkpoint_filename,
monitor=self.monitor_metric,
mode=self.monitor_mode,
save_best_only=not save_every_epoch,
restore_best=not save_every_epoch,
verbose=not save_every_epoch,
temporary_filename=self.best_checkpoint_tmp_filename)
callbacks.append(best_checkpoint)
if save_every_epoch:
best_restore = BestModelRestore(monitor=self.monitor_metric,
mode=self.monitor_mode,
verbose=True)
callbacks.append(best_restore)
if initial_epoch > 1:
# We set the current best metric score in the ModelCheckpoint so that
# it does not save checkpoint it would not have saved if the
# optimization was not stopped.
best_epoch_stats = self.get_best_epoch_stats()
best_epoch = best_epoch_stats['epoch'].item()
best_filename = self.best_checkpoint_filename.format(
epoch=best_epoch)
if not save_every_epoch:
best_checkpoint.best_filename = best_filename
best_checkpoint.current_best = best_epoch_stats[
self.monitor_metric].item()
else:
best_restore.best_weights = torch.load(best_filename,
map_location='cpu')
best_restore.current_best = best_epoch_stats[
self.monitor_metric].item()
return callbacks
def _init_tensorboard_callbacks(self, disable_tensorboard):
tensorboard_writer = None
callbacks = []
if not disable_tensorboard:
if SummaryWriter is None:
warnings.warn(
"tensorboard does not seem to be installed. "
"To remove this warning, set the 'disable_tensorboard' "
"flag to True or install tensorboard.",
stacklevel=3)
else:
tensorboard_writer = SummaryWriter(self.tensorboard_directory)
callbacks += [TensorBoardLogger(tensorboard_writer)]
return tensorboard_writer, callbacks
def _init_lr_scheduler_callbacks(self, lr_schedulers):
callbacks = []
if self.logging:
for i, lr_scheduler in enumerate(lr_schedulers):
filename = self.lr_scheduler_filename % i
tmp_filename = self.lr_scheduler_tmp_filename % i
callbacks += [
LRSchedulerCheckpoint(lr_scheduler,
filename,
verbose=False,
temporary_filename=tmp_filename)
]
else:
callbacks += lr_schedulers
callbacks += [
BestModelRestore(monitor=self.monitor_metric,
mode=self.monitor_mode,
verbose=True)
]
return callbacks
def train(self,
train_generator,
valid_generator=None,
*,
callbacks=None,
lr_schedulers=None,
save_every_epoch=False,
disable_tensorboard=False,
epochs=1000,
steps_per_epoch=None,
validation_steps=None,
batches_per_step=1,
seed=42):
# pylint: disable=too-many-locals
"""
Trains or finetunes the attribute model on a dataset using a generator. If a previous training already occured
and lasted a total of `n_previous` epochs, then the model's weights will be set to the last checkpoint and the
training will be resumed for epochs range (`n_previous`, `epochs`].
If the Experiment has logging enabled (i.e. self.logging is True), numerous callbacks will be automatically
included. Notably, two :class:`~callbacks.ModelCheckpoint` objects will take care of saving the last and every
new best (according to monitor mode) model weights in appropriate checkpoint files.
:class:`~callbacks.OptimizerCheckpoint` and :class:`~callbacks.LRSchedulerCheckpoint` will also respectively
handle the saving of the optimizer and LR scheduler's respective states for future retrieval. Moreover, a
:class:`~callbacks.AtomicCSVLogger` will save all available epoch statistics in an output .tsv file. Lastly, a
:class:`~callbacks.TensorBoardLogger` handles automatic TensorBoard logging of various neural network
statistics.
Args:
train_generator: Generator-like object for the training set. See :func:`~Model.fit_generator()`
for details on the types of generators supported.
valid_generator (optional): Generator-like object for the validation set. See
:func:`~Model.fit_generator()` for details on the types of generators supported.
(Default value = None)
callbacks (List[~poutyne.framework.callbacks.Callback]): List of callbacks that will be called during
training.
(Default value = None)
lr_schedulers (List[~poutyne.framework.callbacks.lr_scheduler._PyTorchLRSchedulerWrapper]): List of
learning rate schedulers.
(Default value = None)
save_every_epoch (bool, optional): Whether or not to save the experiment model's weights after
every epoch.
(Default value = False)
disable_tensorboard (bool, optional): Whether or not to disable the automatic tensorboard logging
callbacks.
(Default value = False)
epochs (int): Number of times the entire training dataset is seen.
(Default value = 1000)
steps_per_epoch (int, optional): Number of batch used during one epoch. Obviously, using this
argument may cause one epoch not to see the entire training dataset or see it multiple times.
(Defaults the number of steps needed to see the entire
training dataset)
validation_steps (int, optional): Same as for ``steps_per_epoch`` but for the validation dataset.
(Defaults to ``steps_per_epoch`` if provided or the number of steps needed to see the entire
validation dataset)
batches_per_step (int): Number of batches on which to compute the running loss before
backpropagating it through the network. Note that the total loss used for backpropagation is
the mean of the `batches_per_step` batch losses.
(Default value = 1)
seed (int, optional): Seed used to make the sampling deterministic.
(Default value = 42)
Returns:
List of dict containing the history of each epoch.
"""
set_seeds(seed)
callbacks = [] if callbacks is None else callbacks
lr_schedulers = [] if lr_schedulers is None else lr_schedulers
# Copy callback list.
callbacks = list(callbacks)
tensorboard_writer = None
initial_epoch = 1
if self.logging:
if not os.path.exists(self.directory):
os.makedirs(self.directory)
# Restarting optimization if needed.
initial_epoch = self._load_epoch_state(lr_schedulers)
callbacks += [
AtomicCSVLogger(self.log_filename,
separator='\t',
append=initial_epoch != 1,
temporary_filename=self.log_tmp_filename)
]
callbacks += self._init_model_restoring_callbacks(
initial_epoch, save_every_epoch)
callbacks += [
ModelCheckpoint(
self.model_checkpoint_filename,
verbose=False,
temporary_filename=self.model_checkpoint_tmp_filename)
]
callbacks += [
OptimizerCheckpoint(
self.optimizer_checkpoint_filename,
verbose=False,
temporary_filename=self.optimizer_checkpoint_tmp_filename)
]
# We save the last epoch number after the end of the epoch so that the
# _load_epoch_state() knows which epoch to restart the optimization.
callbacks += [
PeriodicSaveLambda(
lambda fd, epoch, logs: print(epoch, file=fd),
self.epoch_filename,
temporary_filename=self.epoch_tmp_filename,
open_mode='w')
]
tensorboard_writer, cb_list = self._init_tensorboard_callbacks(
disable_tensorboard)
callbacks += cb_list
# This method returns callbacks that checkpoints the LR scheduler if logging is enabled.
# Otherwise, it just returns the list of LR schedulers with a BestModelRestore callback.
callbacks += self._init_lr_scheduler_callbacks(lr_schedulers)
try:
return self.model.fit_generator(train_generator,
valid_generator,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
validation_steps=validation_steps,
batches_per_step=batches_per_step,
initial_epoch=initial_epoch,
callbacks=callbacks)
finally:
if tensorboard_writer is not None:
tensorboard_writer.close()
def load_checkpoint(self, checkpoint, *, verbose=False):
"""
Loads the attribute model's weights with the weights at a given checkpoint epoch.
Args:
checkpoint (Union[int, str]): Which checkpoint to load the model's weights form.
If 'best', will load the best weights according to ``monitor_metric`` and ``monitor_mode``.
If 'last', will load the last model checkpoint. If int, will load the checkpoint of the
specified epoch.
verbose (bool, optional): Whether or not to print the best epoch number and stats when
checkpoint is 'best'.
(Default value = False)
Returns:
If checkpoint is 'best', will return the best epoch stats, as per :func:`~get_best_epoch_stats()`,
else None.
"""
best_epoch_stats = None
if isinstance(checkpoint, int):
self._load_epoch_checkpoint(checkpoint)
elif checkpoint == 'best':
best_epoch_stats = self._load_best_checkpoint(verbose=verbose)
elif checkpoint == 'last':
self._load_last_checkpoint()
else:
raise ValueError(
"checkpoint argument must be either 'best', 'last' or int. Found : {}"
.format(checkpoint))
return best_epoch_stats
def _load_epoch_checkpoint(self, epoch):
ckpt_filename = self.best_checkpoint_filename.format(epoch=epoch)
if not os.path.isfile(ckpt_filename):
raise ValueError("No checkpoint found for epoch {}".format(epoch))
self.model.load_weights(ckpt_filename)
def _load_best_checkpoint(self, *, verbose=False):
best_epoch_stats = self.get_best_epoch_stats()
best_epoch = best_epoch_stats['epoch'].item()
if verbose:
metrics_str = ', '.join(
'%s: %g' % (metric_name, best_epoch_stats[metric_name].item())
for metric_name in best_epoch_stats.columns[2:])
print("Found best checkpoint at epoch: {}".format(best_epoch))
print(metrics_str)
self._load_epoch_checkpoint(best_epoch)
return best_epoch_stats
def _load_last_checkpoint(self):
self.model.load_weights(self.model_checkpoint_filename)
def test(self,
test_generator,
*,
callbacks=None,
steps=None,
checkpoint='best',
seed=42):
"""
Computes and returns the loss and the metrics of the attribute model on a given test examples
generator.
If the Experiment has logging enabled (i.e. self.logging is True), test and validation statistics
are saved in a specific test output .tsv file.
Args:
test_generator: Generator-like object for the test set. See :func:`~Model.fit_generator()` for
details on the types of generators supported.
callbacks (List[~poutyne.framework.callbacks.Callback]): List of callbacks that will be called during
the testing.
(Default value = None)
steps (int, optional): Number of iterations done on ``generator``.
(Defaults the number of steps needed to see the entire dataset)
checkpoint (Union[str, int]): Which model checkpoint weights to load for the test evaluation.
If 'best', will load the best weights according to ``monitor_metric`` and ``monitor_mode``.
If 'last', will load the last model checkpoint. If int, will load the checkpoint of the
specified epoch.
(Default value = 'best')
seed (int, optional): Seed used to make the sampling deterministic.
(Default value = 42)
If the Experiment has logging enabled (i.e. self.logging is True), one callback will be automatically
included to saved the test metrics. Moreover, a :class:`~callbacks.AtomicCSVLogger` will save the test
metrics in an output .tsv file.
Returns:
dict sorting of all the test metrics values by their names.
"""
set_seeds(seed)
callbacks = [] if callbacks is None else callbacks
# Copy callback list.
callbacks = list(callbacks)
best_epoch_stats = self.load_checkpoint(checkpoint)
if len(self.model.metrics_names) > 0:
test_loss, test_metrics = self.model.evaluate_generator(
test_generator, steps=steps, callbacks=callbacks)
if not isinstance(test_metrics, np.ndarray):
test_metrics = np.array([test_metrics])
else:
test_loss = self.model.evaluate_generator(test_generator,
steps=steps,
callbacks=callbacks)
test_metrics = np.array([])
test_metrics_names = ['test_loss'] + \
['test_' + metric_name for metric_name in self.model.metrics_names]
test_metrics_values = np.concatenate(([test_loss], test_metrics))
test_metrics_dict = dict(zip(test_metrics_names, test_metrics_values))
test_metrics_str = ', '.join('%s: %g' % (col, val)
for col, val in test_metrics_dict.items())
print("On best model: %s" % test_metrics_str)
if self.logging:
test_stats = pd.DataFrame([test_metrics_values],
columns=test_metrics_names)
if best_epoch_stats is not None:
best_epoch_stats = best_epoch_stats.reset_index(drop=True)
test_stats = best_epoch_stats.join(test_stats)
test_stats.to_csv(self.test_log_filename, sep='\t', index=False)
return test_metrics_dict
class ModelCheckpointTest(TestCase):
batch_size = 20
def setUp(self):
torch.manual_seed(42)
self.pytorch_module = nn.Linear(1, 1)
self.loss_function = nn.MSELoss()
self.optimizer = torch.optim.SGD(self.pytorch_module.parameters(),
lr=1e-3)
self.model = Model(self.pytorch_module, 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_batch_begin(1, {})
loss = self._update_model(generator)
checkpointer.on_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_module.zero_grad()
x, y = next(generator)
pred_y = self.pytorch_module(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 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
}
self._test_callbacks_train(params, logs)
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
}
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))
from poutyne.framework import Model
model = Model(network, optimizer, loss_function)
model.to(device)
model.fit_generator(train_loader,
valid_loader,
epochs=num_epochs,
callbacks=callbacks)
test_loss = model.evaluate_generator(test_loader)
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) 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])
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}
self._test_callbacks_train(params, logs)
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 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_batch_begin(1, {})
loss = self._update_model(generator)
earlystopper.on_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)
project_name="dl-pytorch-template",
workspace="francescosaveriozuppichini")
experiment.log_parameters(params)
# create our special resnet18
cnn = resnet18(2).to(device)
# print the model summary to show useful information
logging.info(summary(cnn, (3, 224, 244)))
# define custom optimizer and instantiace the trainer `Model`
optimizer = optim.Adam(cnn.parameters(), lr=params['lr'])
model = Model(cnn, optimizer, "cross_entropy",
batch_metrics=["accuracy"]).to(device)
# usually you want to reduce the lr on plateau and store the best model
callbacks = [
ReduceLROnPlateau(monitor="val_acc", patience=5, verbose=True),
ModelCheckpoint(str(project.checkpoint_dir / f"{time.time()}-model.pt"),
save_best_only="True",
verbose=True),
EarlyStopping(monitor="val_acc", patience=10, mode='max'),
CometCallback(experiment)
]
model.fit_generator(
train_dl,
val_dl,
epochs=50,
callbacks=callbacks,
)
# get the results on the test set
loss, test_acc = model.evaluate_generator(test_dl)
logging.info(f'test_acc=({test_acc})')
experiment.log_metric('test_acc', test_acc)
def train(args):
set_random_seed(42)
model = get_model(args.network, args.classification_head)
print('Loading model')
model.encoder.conv1 = nn.Conv2d(count_channels(args.channels) *
args.neighbours,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
model, device = UtilsFactory.prepare_model(model)
train_df = pd.read_csv(args.train_df).to_dict('records')
val_df = pd.read_csv(args.val_df).to_dict('records')
ds = Dataset(args.channels, args.dataset_path, args.image_size,
args.batch_size, args.num_workers, args.neighbours,
args.classification_head)
loaders = ds.create_loaders(train_df, val_df)
save_path = os.path.join(args.logdir, args.name)
optimizer = get_optimizer(args.optimizer, args.lr, model)
if not args.classification_head:
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=[10, 40, 80, 150, 300], gamma=0.1)
criterion = get_loss(args.loss)
runner = SupervisedRunner()
if args.model_weights_path:
checkpoint = torch.load(args.model_weights_path,
map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
runner.train(model=model,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
loaders=loaders,
callbacks=[DiceCallback()],
logdir=save_path,
num_epochs=args.epochs,
verbose=True)
infer_loader = collections.OrderedDict([('infer', loaders['valid'])])
runner.infer(
model=model,
loaders=infer_loader,
callbacks=[
CheckpointCallback(resume=f'{save_path}/checkpoints/best.pth'),
InferCallback()
],
)
else:
criterion = get_loss('multi')
net = Model(model,
optimizer,
criterion,
batch_metrics=[
classification_head_accuracy, segmentation_head_dice
])
net = net.to(device)
net.fit_generator(loaders['train'],
loaders['valid'],
epochs=args.epochs,
callbacks=[
ModelCheckpoint(
f'{save_path}/checkpoints/best.pth', ),
MultiStepLR(milestones=[10, 40, 80, 150, 300],
gamma=0.1)
])
def training_loop(model,
loss_function,
metrics,
optimizer,
meta_data,
config,
save_path,
train,
valid,
steps_per_epoch,
custom_callbacks=[],
checkpoint_monitor="val_acc",
use_tb=False,
reload=True,
n_epochs=100,
save_freq=1,
save_history_every_k_examples=-1):
callbacks = list(custom_callbacks)
if reload:
H, epoch_start = _reload(model, optimizer, save_path, callbacks)
else:
save_weights(model, optimizer,
os.path.join(save_path, "init_weights.pt"))
history_csv_path, history_pkl_path = os.path.join(
save_path, "history.csv"), os.path.join(save_path, "history.pkl")
logger.info("Removing {} and {}".format(history_pkl_path,
history_csv_path))
os.system("rm " + history_pkl_path)
os.system("rm " + history_csv_path)
H, epoch_start = {}, 0
callbacks += _construct_default_callbacks(model, optimizer, H, save_path,
checkpoint_monitor, save_freq,
custom_callbacks, use_tb,
save_history_every_k_examples)
# Configure callbacks
for clbk in callbacks:
clbk.set_save_path(save_path)
clbk.set_model(model, ignore=False) # TODO: Remove this trick
clbk.set_optimizer(optimizer)
clbk.set_meta_data(meta_data)
clbk.set_config(config)
model = Model(model=model,
optimizer=optimizer,
loss_function=loss_function,
metrics=metrics)
if torch.cuda.is_available():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
logger.info("Sending model to {}".format(device))
model.to(device)
_ = model.fit_generator(
train,
initial_epoch=epoch_start,
steps_per_epoch=steps_per_epoch,
epochs=n_epochs - 1, # Weird convention
verbose=1,
valid_generator=valid,
callbacks=callbacks)
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.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)
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
}
self._test_callbacks_train(params, logs)
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 = list(self.model.get_weight_copies().values())
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 = list(self.model.get_weight_copies().values())
np.testing.assert_almost_equal(returned_params,
expected_params,
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 = list(self.model.get_weight_copies().values())
self.model.set_weights(initial_params)
self.model.fit_generator(list(zip(x, y)),
None,
epochs=1,
batches_per_step=2)
returned_params = list(self.model.get_weight_copies().values())
np.testing.assert_almost_equal(returned_params,
expected_params,
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 = list(self.model.get_weight_copies().values())
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 = list(self.model.get_weight_copies().values())
np.testing.assert_almost_equal(returned_params,
expected_params,
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
}
self._test_callbacks_train(params, logs)
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
}
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)
class IterableMock:
def __init__(self, iterable):
self.iterable = iterable
self.iter = None
self.calls = []
def __iter__(self):
self.calls.append('__iter__')
self.iter = iter(self.iterable)
return self
def __next__(self):
self.calls.append('__next__')
return next(self.iter)
def __len__(self):
self.calls.append('__len__')
return len(self.iterable)
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
}
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
}
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_generator = SomeDataGeneratorWithLen(batch_size=15,
length=10,
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
}
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_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_callback(self):
num_steps = 10
generator = some_data_tensor_generator(ModelTest.batch_size)
result_log = self.model.evaluate_generator(
generator,
steps=num_steps,
return_pred=True,
callbacks=[self.mock_callback])
params = {'batch': ModelTest.epochs}
self._test_callbacks_test(params, result_log)
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)
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_device(self, device):
for p in self.pytorch_network.parameters():
self.assertEqual(p.device, device)
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 OptimizerCheckpointTest(TestCase):
batch_size = 20
epochs = 10
def setUp(self):
torch.manual_seed(42)
self.pytorch_module = nn.Linear(1, 1)
self.loss_function = nn.MSELoss()
self.optimizer = torch.optim.Adam(self.pytorch_module.parameters(),
lr=1e-3)
self.model = Model(self.pytorch_module, 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_batch_begin(1, {})
loss = self._update_model(generator)
checkpointer.on_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_module.zero_grad()
x, y = next(generator)
pred_y = self.pytorch_module(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)
