def test_external_schedulers_provider_hf_transformers(tmpdir, optim, sched, use_datamodule, limit): model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.LogSoftmax()) task = ClassificationTask(model, optimizer=deepcopy(optim), lr_scheduler=deepcopy(sched), loss_fn=F.nll_loss) if limit is not None: batch_count = limit if isinstance(limit, int) else int(limit * 10) trainer = flash.Trainer(max_epochs=1, limit_train_batches=limit) else: batch_count = 10 trainer = flash.Trainer(max_epochs=1) ds = DummyDataset(num_samples=10) if use_datamodule: class TestDataModule(LightningDataModule): def train_dataloader(self): return DataLoader(ds) trainer.fit(task, datamodule=TestDataModule()) else: trainer.fit(task, train_dataloader=DataLoader(ds)) assert task.get_num_training_steps() == batch_count assert isinstance(trainer.optimizers[0], torch.optim.Adadelta) assert isinstance(trainer.lr_schedulers[0]["scheduler"], torch.optim.lr_scheduler.LambdaLR)
def test_no_validation_loop(simple_datamodule): active_learning_dm = ActiveLearningDataModule( simple_datamodule, initial_num_labels=2, query_size=100, val_split=0.0, ) assert active_learning_dm.val_dataloader is None head = nn.Sequential( nn.Dropout(p=0.1), nn.Linear(512, active_learning_dm.num_classes), ) model = ImageClassifier(backbone="resnet18", head=head, num_classes=active_learning_dm.num_classes) trainer = flash.Trainer(max_epochs=3) active_learning_loop = ActiveLearningLoop(label_epoch_frequency=1, inference_iteration=3) active_learning_loop.connect(trainer.fit_loop) trainer.fit_loop = active_learning_loop # Check that we can finetune without val_set trainer.finetune(model, datamodule=active_learning_dm, strategy="no_freeze")
def main(args): seed_everything(args.seed) gpus = 1 if torch.cuda.is_available() else 0 active_dm: ActiveLearningDataModule = get_data_module( args.heuristic, args.data_path) model: ImageClassifier = get_model(active_dm.labelled) logger = TensorBoardLogger( os.path.join(args.ckpt_path, "tensorboard"), name=f"flash-example-cifar-{args.heuristic}-{args.seed}", ) # We use Flash trainer without validation set. # In practice, using a validation set is risky because we overfit often. trainer = flash.Trainer( gpus=gpus, max_epochs=2500, default_root_dir=args.ckpt_path, logger=logger, limit_val_batches=0, ) # We will train for 20 epochs before doing 20 MC-Dropout iterations to estimate uncertainty. active_learning_loop = ActiveLearningLoop(label_epoch_frequency=20, inference_iteration=20) active_learning_loop.connect(trainer.fit_loop) trainer.fit_loop = active_learning_loop # We do not freeze the backbone, this gives better performance. trainer.finetune(model, datamodule=active_dm, strategy="no_freeze")
def test_classification_task_metrics(): train_dataset = FixedDataset([0, 1]) val_dataset = FixedDataset([1, 1]) test_dataset = FixedDataset([0, 0]) model = OnesModel() class CheckAccuracy(Callback): def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: assert math.isclose( trainer.callback_metrics["train_accuracy_epoch"], 0.5) def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: assert math.isclose(trainer.callback_metrics["val_accuracy"], 1.0) def on_test_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: assert math.isclose(trainer.callback_metrics["test_accuracy"], 0.0) task = ClassificationTask(model) trainer = flash.Trainer(max_epochs=1, callbacks=CheckAccuracy(), gpus=torch.cuda.device_count()) trainer.fit(task, train_dataloader=DataLoader(train_dataset), val_dataloaders=DataLoader(val_dataset)) trainer.test(task, DataLoader(test_dataset))
def test_optimizers_and_schedulers(tmpdir, optim, sched, interval): model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.LogSoftmax()) task = ClassificationTask(model, optimizer=optim, lr_scheduler=sched) train_dl = torch.utils.data.DataLoader(DummyDataset()) if sched is None: optimizer = task.configure_optimizers() assert isinstance(optimizer, torch.optim.Adadelta) else: optimizer, scheduler = task.configure_optimizers() assert isinstance(optimizer[0], torch.optim.Adadelta) scheduler = scheduler[0] assert isinstance(scheduler["scheduler"], torch.optim.lr_scheduler.StepLR) assert scheduler["interval"] == interval # generate a checkpoint trainer = flash.Trainer( default_root_dir=tmpdir, limit_train_batches=10, max_epochs=1, ) trainer.fit(task, train_dl)
def test_fast_dev_run_smoke(sample_data): """Test that fast dev run works with the NBeats example data.""" data, training_cutoff, max_prediction_length = sample_data datamodule = TabularForecastingData.from_data_frame( time_idx="time_idx", target="value", categorical_encoders={"series": NaNLabelEncoder().fit(data.series)}, group_ids=["series"], time_varying_unknown_reals=["value"], max_encoder_length=60, max_prediction_length=max_prediction_length, train_data_frame=data[lambda x: x.time_idx <= training_cutoff], val_data_frame=data, batch_size=4, ) model = TabularForecaster( datamodule.parameters, backbone="n_beats", backbone_kwargs={ "widths": [32, 512], "backcast_loss_ratio": 0.1 }, ) trainer = flash.Trainer(max_epochs=1, fast_dev_run=True, gradient_clip_val=0.01) trainer.fit(model, datamodule=datamodule)
def test_vissl_training(backbone, training_strategy, head, pretraining_transform, embedding_size): datamodule = ImageClassificationData.from_datasets( train_dataset=FakeData(16), predict_dataset=FakeData(8), batch_size=4, ) embedder = ImageEmbedder( backbone=backbone, training_strategy=training_strategy, head=head, pretraining_transform=pretraining_transform, ) trainer = flash.Trainer( max_steps=3, max_epochs=1, gpus=torch.cuda.device_count(), ) trainer.fit(embedder, datamodule=datamodule) predictions = trainer.predict(embedder, datamodule=datamodule) for prediction_batch in predictions: for prediction in prediction_batch: assert prediction.size(0) == embedding_size
def test_testing_raises(sample_data): """Tests that ``NotImplementedError`` is raised when attempting to perform a test pass.""" data, training_cutoff, max_prediction_length = sample_data datamodule = TabularForecastingData.from_data_frame( time_idx="time_idx", target="value", categorical_encoders={"series": NaNLabelEncoder().fit(data.series)}, group_ids=["series"], time_varying_unknown_reals=["value"], max_encoder_length=60, max_prediction_length=max_prediction_length, train_data_frame=data[lambda x: x.time_idx <= training_cutoff], test_data_frame=data, batch_size=4, ) model = TabularForecaster( datamodule.parameters, backbone="n_beats", backbone_kwargs={ "widths": [32, 512], "backcast_loss_ratio": 0.1 }, ) trainer = flash.Trainer(max_epochs=1, fast_dev_run=True, gradient_clip_val=0.01) with pytest.raises( NotImplementedError, match= "Backbones provided by PyTorch Forecasting don't support testing." ): trainer.test(model, datamodule=datamodule)
def test_video_classifier_finetune_from_csv(tmpdir): with mock_video_csv_file(tmpdir) as (mock_csv, total_duration): half_duration = total_duration / 2 - 1e-9 datamodule = VideoClassificationData.from_csv( "file", "target", train_file=mock_csv, clip_sampler="uniform", clip_duration=half_duration, video_sampler=SequentialSampler, decode_audio=False, batch_size=1, ) for sample in datamodule.train_dataset.data: expected_t_shape = 5 assert sample["video"].shape[1] == expected_t_shape model = VideoClassifier(num_classes=datamodule.num_classes, pretrained=False, backbone="slow_r50") trainer = flash.Trainer(default_root_dir=tmpdir, fast_dev_run=True, gpus=torch.cuda.device_count()) trainer.finetune(model, datamodule=datamodule)
def test_video_classifier_finetune_fiftyone(tmpdir): with mock_encoded_video_dataset_folder(tmpdir) as ( dir_name, total_duration, ): half_duration = total_duration / 2 - 1e-9 train_dataset = fo.Dataset.from_dir( dir_name, dataset_type=fo.types.VideoClassificationDirectoryTree, ) datamodule = VideoClassificationData.from_fiftyone( train_dataset=train_dataset, clip_sampler="uniform", clip_duration=half_duration, video_sampler=SequentialSampler, decode_audio=False, batch_size=1, ) for sample in datamodule.train_dataset.data: expected_t_shape = 5 assert sample["video"].shape[1] == expected_t_shape model = VideoClassifier(num_classes=datamodule.num_classes, pretrained=False, backbone="slow_r50") trainer = flash.Trainer(fast_dev_run=True, gpus=torch.cuda.device_count()) trainer.finetune(model, datamodule=datamodule)
def test_optimization(tmpdir): model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.LogSoftmax()) optim = torch.optim.Adam(model.parameters()) task = ClassificationTask(model, optimizer=optim, scheduler=None) optimizer = task.configure_optimizers() assert optimizer == optim task = ClassificationTask(model, optimizer=torch.optim.Adadelta, optimizer_kwargs={"eps": 0.5}, scheduler=None) optimizer = task.configure_optimizers() assert isinstance(optimizer, torch.optim.Adadelta) assert optimizer.defaults["eps"] == 0.5 task = ClassificationTask( model, optimizer=torch.optim.Adadelta, scheduler=torch.optim.lr_scheduler.StepLR, scheduler_kwargs={"step_size": 1} ) optimizer, scheduler = task.configure_optimizers() assert isinstance(optimizer[0], torch.optim.Adadelta) assert isinstance(scheduler[0], torch.optim.lr_scheduler.StepLR) optim = torch.optim.Adadelta(model.parameters()) task = ClassificationTask(model, optimizer=optim, scheduler=torch.optim.lr_scheduler.StepLR(optim, step_size=1)) optimizer, scheduler = task.configure_optimizers() assert isinstance(optimizer[0], torch.optim.Adadelta) assert isinstance(scheduler[0], torch.optim.lr_scheduler.StepLR) if _TRANSFORMERS_AVAILABLE: from transformers.optimization import get_linear_schedule_with_warmup assert task.available_schedulers() == [ 'constant_schedule', 'constant_schedule_with_warmup', 'cosine_schedule_with_warmup', 'cosine_with_hard_restarts_schedule_with_warmup', 'linear_schedule_with_warmup', 'polynomial_decay_schedule_with_warmup' ] optim = torch.optim.Adadelta(model.parameters()) with pytest.raises(MisconfigurationException, match="The LightningModule isn't attached to the trainer yet."): task = ClassificationTask(model, optimizer=optim, scheduler="linear_schedule_with_warmup") optimizer, scheduler = task.configure_optimizers() task = ClassificationTask( model, optimizer=optim, scheduler="linear_schedule_with_warmup", scheduler_kwargs={"num_warmup_steps": 0.1}, loss_fn=F.nll_loss, ) trainer = flash.Trainer(max_epochs=1, limit_train_batches=2) ds = DummyDataset() trainer.fit(task, train_dataloader=DataLoader(ds)) optimizer, scheduler = task.configure_optimizers() assert isinstance(optimizer[0], torch.optim.Adadelta) assert isinstance(scheduler[0], torch.optim.lr_scheduler.LambdaLR) expected = get_linear_schedule_with_warmup.__name__ assert scheduler[0].lr_lambdas[0].__qualname__.split('.')[0] == expected
def test_finetuning_errors_and_exceptions(strategy): task = TestTaskWithFinetuning(loss_fn=F.nll_loss) trainer = flash.Trainer(max_epochs=1, limit_train_batches=10) ds = DummyDataset() with pytest.raises(MisconfigurationException): trainer.finetune(task, train_dataloader=DataLoader(ds), strategy=strategy)
def test_finetuning(tmpdir, strategy, checker_class, checker_class_data): task = TestTaskWithFinetuning(loss_fn=F.nll_loss) callbacks = [] if checker_class is None else checker_class( dirpath=tmpdir, **checker_class_data) trainer = flash.Trainer(max_epochs=5, limit_train_batches=10, callbacks=callbacks) ds = DummyDataset() trainer.finetune(task, train_dataloader=DataLoader(ds), strategy=strategy)
def test_fastface_training(): dataset = ff.dataset.FDDBDataset(source_dir="data/", phase="val") datamodule = FaceDetectionData.from_datasets(train_dataset=dataset, predict_dataset=dataset, batch_size=2) model = FaceDetector(model="lffd_slim") # test fit trainer = flash.Trainer(fast_dev_run=True) trainer.finetune(model, datamodule=datamodule, strategy="freeze") trainer.predict(model, datamodule=datamodule)
def _test_fit(self, tmpdir, task_kwargs): """Tests that a single batch fit pass completes.""" dataset = _StaticDataset(self.example_train_sample, 4) args = self.task_args kwargs = dict(**self.task_kwargs) kwargs.update(task_kwargs) model = self.task(*args, **kwargs) trainer = flash.Trainer(default_root_dir=tmpdir, fast_dev_run=True) trainer.fit(model, model.process_train_dataset(dataset, batch_size=4))
def test_predict(tmpdir): datamodule = TextClassificationData.from_lists(predict_data=predict_data, batch_size=4) model = TextEmbedder(backbone=TEST_BACKBONE) trainer = flash.Trainer(gpus=torch.cuda.device_count()) predictions = trainer.predict(model, datamodule=datamodule) assert [t.size() for t in predictions[0] ] == [torch.Size([384]), torch.Size([384]), torch.Size([384])]
def test_only_embedding(backbone, embedding_size): datamodule = ImageClassificationData.from_datasets( predict_dataset=FakeData(8), batch_size=4, transform_kwargs=dict(image_size=(224, 224)), ) embedder = ImageEmbedder(backbone=backbone) trainer = flash.Trainer() predictions = trainer.predict(embedder, datamodule=datamodule) for prediction_batch in predictions: for prediction in prediction_batch: assert prediction.size(0) == embedding_size
def test_external_optimizers_torch_optimizer(tmpdir, optim): model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.LogSoftmax()) task = ClassificationTask(model, optimizer=optim, lr_scheduler=None, loss_fn=F.nll_loss) trainer = flash.Trainer(max_epochs=1, limit_train_batches=2, gpus=torch.cuda.device_count()) ds = DummyDataset() trainer.fit(task, train_dataloader=DataLoader(ds)) from torch_optimizer import Yogi optimizer = task.configure_optimizers() assert isinstance(optimizer, Yogi)
def test_detection_fiftyone(tmpdir, head, backbone): train_dataset = _create_synth_fiftyone_dataset(tmpdir) datamodule = ObjectDetectionData.from_fiftyone(train_dataset=train_dataset, batch_size=1) model = ObjectDetector(head=head, backbone=backbone, num_classes=datamodule.num_classes) trainer = flash.Trainer(fast_dev_run=True, gpus=torch.cuda.device_count()) trainer.finetune(model, datamodule, strategy="freeze") test_image_one = os.fspath(tmpdir / "test_one.png") test_image_two = os.fspath(tmpdir / "test_two.png") Image.new("RGB", (512, 512)).save(test_image_one) Image.new("RGB", (512, 512)).save(test_image_two) datamodule = ObjectDetectionData.from_files(predict_files=[str(test_image_one), str(test_image_two)], batch_size=1) trainer.predict(model, datamodule=datamodule)
def test_detection_fiftyone(tmpdir, model, backbone): train_dataset = _create_synth_fiftyone_dataset(tmpdir) data = ObjectDetectionData.from_fiftyone(train_dataset=train_dataset, batch_size=1) model = ObjectDetector(model=model, backbone=backbone, num_classes=data.num_classes) trainer = flash.Trainer(fast_dev_run=True) trainer.finetune(model, data) test_image_one = os.fspath(tmpdir / "test_one.png") test_image_two = os.fspath(tmpdir / "test_two.png") Image.new('RGB', (512, 512)).save(test_image_one) Image.new('RGB', (512, 512)).save(test_image_two) test_images = [str(test_image_one), str(test_image_two)] model.predict(test_images)
def test_detection(tmpdir): train_folder, coco_ann_path = _create_synth_coco_dataset(tmpdir) data = ObjectDetectionData.from_coco(train_folder=train_folder, train_ann_file=coco_ann_path, batch_size=1) model = ObjectDetector(num_classes=data.num_classes) trainer = flash.Trainer(fast_dev_run=True) trainer.finetune(model, data) test_image_one = os.fspath(tmpdir / "test_one.png") test_image_two = os.fspath(tmpdir / "test_two.png") Image.new('RGB', (1920, 1080)).save(test_image_one) Image.new('RGB', (1920, 1080)).save(test_image_two) test_images = [test_image_one, test_image_two] model.predict(test_images)
def run(self, train_folder): # Create a datamodule from the given dataset datamodule = ImageClassificationData.from_folders( train_folder=train_folder, batch_size=1, val_split=0.5, ) # Create an image classfier task with the given backbone model = ImageClassifier(datamodule.num_classes, backbone=self.backbone) # Start a Lightning trainer, with 1 training batch and 4 validation batches trainer = flash.Trainer( max_epochs=self.max_epochs, limit_train_batches=1, limit_val_batches=4, callbacks=[ModelCheckpoint(monitor="val_cross_entropy")], ) # Train the model trainer.fit(model, datamodule=datamodule) # Save the model path self.best_model_path = trainer.checkpoint_callback.best_model_path # Save the model score self.best_model_score = trainer.checkpoint_callback.best_model_score.item( )
def test_detection(tmpdir, head, backbone): train_folder, coco_ann_path = _create_synth_coco_dataset(tmpdir) data = ObjectDetectionData.from_coco(train_folder=train_folder, train_ann_file=coco_ann_path, batch_size=1) model = ObjectDetector(head=head, backbone=backbone, num_classes=data.num_classes) trainer = flash.Trainer(fast_dev_run=True, gpus=torch.cuda.device_count()) trainer.finetune(model, data, strategy="freeze") test_image_one = os.fspath(tmpdir / "test_one.png") test_image_two = os.fspath(tmpdir / "test_two.png") Image.new("RGB", (512, 512)).save(test_image_one) Image.new("RGB", (512, 512)).save(test_image_two) test_images = [str(test_image_one), str(test_image_two)] model.predict(test_images)
os.getcwd(), "data/coco128/annotations/instances_train2017.json")) parser.add_argument('--max_epochs', type=int, default=1) parser.add_argument('--learning_rate', type=float, default=1e-3) parser.add_argument('--gpus', type=int, default=None) args = parser.parse_args() # 1. Download the data if args.download: # Dataset Credit: https://www.kaggle.com/ultralytics/coco128 download_data( "https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0/coco128.zip", os.path.join(os.getcwd(), "data/")) # 2. Load the Data datamodule = ObjectDetectionData.from_coco( train_folder=args.train_folder, train_ann_file=args.train_ann_file, batch_size=2) # 3. Build the model model = ObjectDetector(num_classes=datamodule.num_classes) # 4. Create the trainer trainer = flash.Trainer(max_epochs=args.max_epochs, gpus=args.gpus) # 5. Finetune the model trainer.finetune(model, datamodule) # 6. Save it! trainer.save_checkpoint("object_detection_model.pt")
# 3 Load FiftyOne datasets datamodule = ImageClassificationData.from_fiftyone( train_dataset=train_dataset, val_dataset=val_dataset, test_dataset=test_dataset, ) # 4 Fine tune a model model = ImageClassifier( backbone="resnet18", num_classes=datamodule.num_classes, serializer=Labels(), ) trainer = flash.Trainer( max_epochs=1, limit_train_batches=1, limit_val_batches=1, ) trainer.finetune( model, datamodule=datamodule, strategy=FreezeUnfreeze(unfreeze_epoch=1), ) trainer.save_checkpoint("image_classification_model.pt") # 5 Predict from checkpoint on data with ground truth model = ImageClassifier.load_from_checkpoint("https://flash-weights.s3.amazonaws.com/image_classification_model.pt") model.serializer = FiftyOneLabels(return_filepath=False) datamodule = ImageClassificationData.from_fiftyone(predict_dataset=test_dataset) predictions = trainer.predict(model, datamodule=datamodule)
# limitations under the License. import torch import flash from flash import download_data, Trainer from flash.text import SummarizationData, SummarizationTask # 1. Download the data download_data("https://pl-flash-data.s3.amazonaws.com/xsum.zip", "data/") # 2. Load the data datamodule = SummarizationData.from_files( train_file="data/xsum/train.csv", val_file="data/xsum/valid.csv", test_file="data/xsum/test.csv", input="input", target="target" ) # 3. Build the model model = SummarizationTask() # 4. Create the trainer. Run once on data trainer = flash.Trainer(gpus=int(torch.cuda.is_available()), fast_dev_run=True) # 5. Fine-tune the model trainer.finetune(model, datamodule=datamodule) # 6. Save it! trainer.save_checkpoint("summarization_model_xsum.pt")
# limitations under the License. from flash.text import TextClassificationData, TextClassifier # 1. Download the data download_data("https://pl-flash-data.s3.amazonaws.com/imdb.zip", 'data/') # 2. Load the data datamodule = TextClassificationData.from_files( train_file="data/imdb/train.csv", valid_file="data/imdb/valid.csv", test_file="data/imdb/test.csv", input="review", target="sentiment", batch_size=512 ) # 3. Build the model model = TextClassifier(num_classes=datamodule.num_classes) # 4. Create the trainer. Run once on data trainer = flash.Trainer(max_epochs=1) # 5. Fine-tune the model trainer.finetune(model, datamodule=datamodule, strategy='freeze') # 6. Test model trainer.test() # 7. Save it! trainer.save_checkpoint("text_classification_model.pt")
# 1. Download the data download_data("https://pl-flash-data.s3.amazonaws.com/titanic.zip", "data/") # 2. Load the data datamodule = TabularData.from_csv( ["Sex", "Age", "SibSp", "Parch", "Ticket", "Cabin", "Embarked"], ["Fare"], target_field="Survived", train_file="./data/titanic/titanic.csv", test_file="./data/titanic/test.csv", val_split=0.25, ) # 3. Build the model model = TabularClassifier.from_data( datamodule, metrics=[Accuracy(), Precision(), Recall()]) # 4. Create the trainer trainer = flash.Trainer(fast_dev_run=True) # 5. Train the model trainer.fit(model, datamodule=datamodule) # 6. Test model trainer.test(model) # 7. Save it! trainer.save_checkpoint("tabular_classification_model.pt")
dm = cls.from_load_data_inputs( train_load_data_input=(x_train, y_train), test_load_data_input=(x_test, y_test), preprocess=preprocess, batch_size=batch_size, num_workers=num_workers ) dm.num_inputs = dm.train_dataset.num_inputs return dm x, y = datasets.load_diabetes(return_X_y=True) datamodule = NumpyDataModule.from_dataset(x, y, NumpyPreprocess()) model = RegressionTask(num_inputs=datamodule.num_inputs) trainer = flash.Trainer(max_epochs=10, progress_bar_refresh_rate=20) trainer.fit(model, datamodule=datamodule) predict_data = np.array([[0.0199, 0.0507, 0.1048, 0.0701, -0.0360, -0.0267, -0.0250, -0.0026, 0.0037, 0.0403], [-0.0128, -0.0446, 0.0606, 0.0529, 0.0480, 0.0294, -0.0176, 0.0343, 0.0702, 0.0072], [0.0381, 0.0507, 0.0089, 0.0425, -0.0428, -0.0210, -0.0397, -0.0026, -0.0181, 0.0072], [-0.0128, -0.0446, -0.0235, -0.0401, -0.0167, 0.0046, -0.0176, -0.0026, -0.0385, -0.0384], [-0.0237, -0.0446, 0.0455, 0.0907, -0.0181, -0.0354, 0.0707, -0.0395, -0.0345, -0.0094]]) predictions = model.predict(predict_data) # out: This prediction: tensor([14.7288]) is above the threshold: 14.72 print(predictions) # out: [tensor([14.7190]), tensor([14.7100]), tensor([14.7288]), tensor([14.6685]), tensor([14.6687])]
batch_size=32, ) # 2. Build the task model = TabularForecaster( datamodule.parameters, backbone="n_beats", backbone_kwargs={ "widths": [32, 512], "backcast_loss_ratio": 0.1 }, ) # 3. Create the trainer and train the model trainer = flash.Trainer(max_epochs=1, gpus=torch.cuda.device_count(), gradient_clip_val=0.01) trainer.fit(model, datamodule=datamodule) # 4. Generate predictions datamodule = TabularForecastingData.from_data_frame( predict_data_frame=data, parameters=datamodule.parameters) predictions = trainer.predict(model, datamodule=datamodule) print(predictions) # Plot with PyTorch Forecasting! predictions, inputs = convert_predictions(predictions) fig, axs = plt.subplots(2, 3, sharex="col") for idx in range(3):