Example #1
0
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")
Example #4
0
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))
Example #5
0
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)
Example #6
0
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
Example #8
0
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)
Example #9
0
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)
Example #10
0
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)
Example #11
0
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)
Example #15
0
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))
Example #16
0
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
Example #18
0
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)
Example #20
0
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)
Example #22
0
 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(
     )
Example #23
0
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)
Example #24
0
            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)
Example #26
0
# 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):