Template¶

The Task¶

Here you should add a description of your task. For example: Classification is the task of assigning one of a number of classes to each data point.

Example¶

Note

Here you should add a short intro to your example, and then use literalinclude to add it. To make it simple, you can fill in this template.

Let’s look at the task of <describe the task> using the <data set used in the example>. The dataset contains <describe the data>. Here’s an outline:

<present the folder structure of the data or some data samples here>

Once we’ve downloaded the data using download_data(), we create the <link to the DataModule with :class:>. We select a pre-trained backbone to use for our <link to the Task with :class:> and finetune on the <name of the data set> data. We then use the trained <link to the Task with :class:> for inference. Finally, we save the model. Here’s the full example:

import flash
import numpy as np
import torch
from flash.template import TemplateData, TemplateSKLearnClassifier
from sklearn import datasets

# 1. Create the DataModule
datamodule = TemplateData.from_sklearn(
    train_bunch=datasets.load_iris(),
    val_split=0.1,
    batch_size=4,
)

# 2. Build the task
model = TemplateSKLearnClassifier(num_features=datamodule.num_features, num_classes=datamodule.num_classes)

# 3. Create the trainer and train the model
trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
trainer.fit(model, datamodule=datamodule)

# 4. Classify a few examples
datamodule = TemplateData.from_numpy(
    predict_data=[
        np.array([4.9, 3.0, 1.4, 0.2]),
        np.array([6.9, 3.2, 5.7, 2.3]),
        np.array([7.2, 3.0, 5.8, 1.6]),
    ],
    batch_size=4,
)
predictions = trainer.predict(model, datamodule=datamodule, output="classes")
print(predictions)

# 5. Save the model!
trainer.save_checkpoint("template_model.pt")