def __init__(
self,
input_size: int,
output_size: int,
hidden_size: int,
n_hidden_layers: int,
n_classes: int,
loss=CrossEntropy(),
**kwargs,
):
# saves arguments in signature to `.hparams` attribute, mandatory call - do not skip this
self.save_hyperparameters()
# pass additional arguments to BaseModel.__init__, mandatory call - do not skip this
super().__init__(**kwargs)
self.network = FullyConnectedModule(
input_size=self.hparams.input_size,
output_size=self.hparams.output_size * self.hparams.n_classes,
hidden_size=self.hparams.hidden_size,
n_hidden_layers=self.hparams.n_hidden_layers,
)
def test_integration(multiple_dataloaders_with_covariates, tmp_path, gpus):
train_dataloader = multiple_dataloaders_with_covariates["train"]
val_dataloader = multiple_dataloaders_with_covariates["val"]
early_stop_callback = EarlyStopping(monitor="val_loss",
min_delta=1e-4,
patience=1,
verbose=False,
mode="min")
# check training
logger = TensorBoardLogger(tmp_path)
checkpoint = ModelCheckpoint(filepath=tmp_path)
trainer = pl.Trainer(
checkpoint_callback=checkpoint,
max_epochs=3,
gpus=gpus,
weights_summary="top",
gradient_clip_val=0.1,
callbacks=[early_stop_callback],
fast_dev_run=True,
logger=logger,
)
# test monotone constraints automatically
if "discount_in_percent" in train_dataloader.dataset.reals:
monotone_constaints = {"discount_in_percent": +1}
cuda_context = torch.backends.cudnn.flags(enabled=False)
else:
monotone_constaints = {}
cuda_context = nullcontext()
with cuda_context:
if isinstance(train_dataloader.dataset.target_normalizer,
NaNLabelEncoder):
loss = CrossEntropy()
elif isinstance(train_dataloader.dataset.target_normalizer,
MultiNormalizer):
loss = MultiLoss([
CrossEntropy()
if isinstance(normalizer, NaNLabelEncoder) else QuantileLoss()
for normalizer in
train_dataloader.dataset.target_normalizer.normalizers
])
else:
loss = QuantileLoss()
net = TemporalFusionTransformer.from_dataset(
train_dataloader.dataset,
learning_rate=0.15,
hidden_size=4,
attention_head_size=1,
dropout=0.2,
hidden_continuous_size=2,
loss=loss,
log_interval=5,
log_val_interval=1,
log_gradient_flow=True,
monotone_constaints=monotone_constaints,
)
net.size()
try:
trainer.fit(
net,
train_dataloader=train_dataloader,
val_dataloaders=val_dataloader,
)
# check loading
net = TemporalFusionTransformer.load_from_checkpoint(
checkpoint.best_model_path)
# check prediction
net.predict(val_dataloader,
fast_dev_run=True,
return_index=True,
return_decoder_lengths=True)
# check prediction on gpu
if not (isinstance(gpus, int) and gpus == 0):
net.to("cuda")
net.predict(val_dataloader,
fast_dev_run=True,
return_index=True,
return_decoder_lengths=True)
finally:
shutil.rmtree(tmp_path, ignore_errors=True)
def test_integration(multiple_dataloaders_with_covariates, tmp_path, gpus):
train_dataloader = multiple_dataloaders_with_covariates["train"]
val_dataloader = multiple_dataloaders_with_covariates["val"]
early_stop_callback = EarlyStopping(monitor="val_loss", min_delta=1e-4, patience=1, verbose=False, mode="min")
# check training
logger = TensorBoardLogger(tmp_path)
trainer = pl.Trainer(
max_epochs=2,
gpus=gpus,
weights_summary="top",
gradient_clip_val=0.1,
callbacks=[early_stop_callback],
checkpoint_callback=True,
default_root_dir=tmp_path,
limit_train_batches=2,
limit_val_batches=2,
logger=logger,
)
# test monotone constraints automatically
if "discount_in_percent" in train_dataloader.dataset.reals:
monotone_constaints = {"discount_in_percent": +1}
cuda_context = torch.backends.cudnn.flags(enabled=False)
else:
monotone_constaints = {}
cuda_context = nullcontext()
with cuda_context:
if isinstance(train_dataloader.dataset.target_normalizer, NaNLabelEncoder):
loss = CrossEntropy()
elif isinstance(train_dataloader.dataset.target_normalizer, MultiNormalizer):
loss = MultiLoss(
[
CrossEntropy() if isinstance(normalizer, NaNLabelEncoder) else QuantileLoss()
for normalizer in train_dataloader.dataset.target_normalizer.normalizers
]
)
else:
loss = QuantileLoss()
net = TemporalFusionTransformer.from_dataset(
train_dataloader.dataset,
learning_rate=0.15,
hidden_size=4,
attention_head_size=1,
dropout=0.2,
hidden_continuous_size=2,
loss=loss,
log_interval=5,
log_val_interval=1,
log_gradient_flow=True,
monotone_constaints=monotone_constaints,
)
net.size()
try:
trainer.fit(
net,
train_dataloader=train_dataloader,
val_dataloaders=val_dataloader,
)
# check loading
net = TemporalFusionTransformer.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)
# check prediction
predictions, x, index = net.predict(val_dataloader, return_index=True, return_x=True)
pred_len = len(multiple_dataloaders_with_covariates["val"].dataset)
# check that output is of correct shape
def check(x):
if isinstance(x, (tuple, list)):
for xi in x:
check(xi)
elif isinstance(x, dict):
for xi in x.values():
check(xi)
else:
assert pred_len == x.shape[0], "first dimension should be prediction length"
check(predictions)
check(x)
check(index)
# check prediction on gpu
if not (isinstance(gpus, int) and gpus == 0):
net.to("cuda")
net.predict(val_dataloader, fast_dev_run=True, return_index=True, return_decoder_lengths=True)
finally:
shutil.rmtree(tmp_path, ignore_errors=True)
return_decoder_lengths=True)
@pytest.mark.parametrize(
"kwargs",
[
{},
dict(
loss=MultiLoss([QuantileLoss(), MAE()]),
data_loader_kwargs=dict(
time_varying_unknown_reals=["volume", "discount"],
target=["volume", "discount"],
),
),
dict(
loss=CrossEntropy(),
data_loader_kwargs=dict(target="agency", ),
),
],
)
def test_integration(data_with_covariates, tmp_path, gpus, kwargs):
_integration(data_with_covariates.assign(target=lambda x: x.volume),
tmp_path, gpus, **kwargs)
@pytest.fixture
def model(dataloaders_with_covariates):
dataset = dataloaders_with_covariates["train"].dataset
net = DecoderMLP.from_dataset(
dataset,
learning_rate=0.15,