def test_jit(tmpdir, jitter, args):
path = os.path.join(tmpdir, "test.pt")
model = ImageClassifier(2)
model.eval()
model = jitter(model, *args)
torch.jit.save(model, path)
model = torch.jit.load(path)
out = model(torch.rand(1, 3, 32, 32))
assert isinstance(out, torch.Tensor)
assert out.shape == torch.Size([1, 2])
def test_multilabel(tmpdir):
num_classes = 4
ds = DummyMultiLabelDataset(num_classes)
model = ImageClassifier(num_classes, multi_label=True, serializer=Probabilities(multi_label=True))
train_dl = torch.utils.data.DataLoader(ds, batch_size=2)
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
trainer.finetune(model, train_dl, strategy="freeze_unfreeze")
image, label = ds[0][DefaultDataKeys.INPUT], ds[0][DefaultDataKeys.TARGET]
predictions = model.predict([{DefaultDataKeys.INPUT: image}])
assert (torch.tensor(predictions) > 1).sum() == 0
assert (torch.tensor(predictions) < 0).sum() == 0
assert len(predictions[0]) == num_classes == len(label)
assert len(torch.unique(label)) <= 2
def test_default_strategies(tmpdir):
num_classes = 10
ds = DummyDataset()
model = ImageClassifier(num_classes, backbone="resnet50")
trainer = Trainer(fast_dev_run=2)
trainer.fit(model, train_dataloader=DataLoader(ds))
def test_wrongly_specified_training_strategies():
with pytest.raises(KeyError, match="something is not in FlashRegistry"):
ImageClassifier(
backbone="resnet18",
training_strategy="something",
training_strategy_kwargs={"ways": 2, "shots": 4, "meta_batch_size": 10},
)
def test_classification_fiftyone(tmpdir):
tmpdir = Path(tmpdir)
(tmpdir / "a").mkdir()
(tmpdir / "b").mkdir()
_rand_image().save(tmpdir / "a_1.png")
_rand_image().save(tmpdir / "b_1.png")
train_images = [
str(tmpdir / "a_1.png"),
str(tmpdir / "b_1.png"),
]
train_dataset = fo.Dataset.from_dir(str(tmpdir),
dataset_type=fo.types.ImageDirectory)
s1 = train_dataset[train_images[0]]
s2 = train_dataset[train_images[1]]
s1["test"] = fo.Classification(label="1")
s2["test"] = fo.Classification(label="2")
s1.save()
s2.save()
data = ImageClassificationData.from_fiftyone(
train_dataset=train_dataset,
label_field="test",
batch_size=2,
num_workers=0,
image_size=(64, 64),
)
model = ImageClassifier(num_classes=2, backbone="resnet18")
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
trainer.finetune(model, datamodule=data, strategy="freeze")
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 _test_learn2learning_training_strategies(gpus, accelerator, training_strategy, tmpdir):
train_dir = Path(tmpdir / "train")
train_dir.mkdir()
(train_dir / "a").mkdir()
pa_1 = train_dir / "a" / "1.png"
pa_2 = train_dir / "a" / "2.png"
pb_1 = train_dir / "b" / "1.png"
pb_2 = train_dir / "b" / "2.png"
image_size = (96, 96)
_rand_image(image_size).save(pa_1)
_rand_image(image_size).save(pa_2)
(train_dir / "b").mkdir()
_rand_image(image_size).save(pb_1)
_rand_image(image_size).save(pb_2)
n = 5
dm = ImageClassificationData.from_files(
train_files=[str(pa_1)] * n + [str(pa_2)] * n + [str(pb_1)] * n + [str(pb_2)] * n,
train_targets=[0] * n + [1] * n + [2] * n + [3] * n,
batch_size=1,
num_workers=0,
transform_kwargs=dict(image_size=image_size),
)
model = ImageClassifier(
backbone="resnet18",
training_strategy=training_strategy,
training_strategy_kwargs={"ways": dm.num_classes, "shots": 4, "meta_batch_size": 4},
)
trainer = Trainer(fast_dev_run=2, gpus=gpus, accelerator=accelerator)
trainer.fit(model, datamodule=dm)
def test_available_backbones():
backbones = ImageClassifier.available_backbones()
assert "resnet152" in backbones
class Foo(ImageClassifier):
backbones = None
assert Foo.available_backbones() == {}
def test_classification_task_predict_folder_path(tmpdir):
train_dir = Path(tmpdir / "train")
train_dir.mkdir()
def _rand_image():
return Image.fromarray(
np.random.randint(0, 255, (256, 256, 3), dtype="uint8"))
_rand_image().save(train_dir / "1.png")
_rand_image().save(train_dir / "2.png")
datamodule = ImageClassificationData.from_folders(predict_folder=train_dir)
task = ImageClassifier(num_classes=10)
predictions = task.predict(str(train_dir),
data_pipeline=datamodule.data_pipeline)
assert len(predictions) == 2
def test_classification(tmpdir):
tmpdir = Path(tmpdir)
(tmpdir / "a").mkdir()
(tmpdir / "b").mkdir()
image_a = str(tmpdir / "a" / "a_1.png")
image_b = str(tmpdir / "b" / "b_1.png")
_rand_image().save(image_a)
_rand_image().save(image_b)
data = ImageClassificationData.from_files(
train_files=[image_a, image_b],
train_targets=[0, 1],
num_workers=0,
batch_size=2,
image_size=(64, 64),
)
model = ImageClassifier(num_classes=2, backbone="resnet18")
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
trainer.finetune(model, datamodule=data, strategy="freeze")
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 get_model(dm):
loss_fn = nn.CrossEntropyLoss()
head = nn.Sequential(
nn.Linear(512, 512),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(512, 512),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(512, dm.num_classes),
)
LR = 0.001
model = ImageClassifier(
num_classes=dm.num_classes,
head=head,
backbone="vgg16",
pretrained=True,
loss_fn=loss_fn,
optimizer=partial(torch.optim.SGD, momentum=0.9, weight_decay=5e-4),
learning_rate=LR,
serializer=Logits(
), # Note the serializer to Logits to be able to estimate uncertainty.
)
return model
def test_available_backbones_raises():
with pytest.raises(ModuleNotFoundError,
match="Required dependencies not available."):
_ = ImageClassifier.available_backbones()
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from flash import Trainer
from flash.core.classification import Probabilities
from flash.core.data.utils import download_data
from flash.image import ImageClassificationData, ImageClassifier
# 1. Download the data
download_data("https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip",
"data/")
# 2. Load the model from a checkpoint
model = ImageClassifier.load_from_checkpoint(
"https://flash-weights.s3.amazonaws.com/image_classification_model.pt")
# 3a. Predict what's on a few images! ants or bees?
model.serializer = Probabilities()
predictions = model.predict([
"data/hymenoptera_data/val/bees/65038344_52a45d090d.jpg",
"data/hymenoptera_data/val/bees/590318879_68cf112861.jpg",
"data/hymenoptera_data/val/ants/540543309_ddbb193ee5.jpg",
])
print(predictions)
# 3b. Or generate predictions with a whole folder!
datamodule = ImageClassificationData.from_folders(
predict_folder="data/hymenoptera_data/predict/")
def test_unfreeze():
model = ImageClassifier(2)
model.unfreeze()
for p in model.backbone.parameters():
assert p.requires_grad is True
def test_non_existent_backbone():
with pytest.raises(KeyError):
ImageClassifier(2, "i am never going to implement this lol")
def test_init_train(tmpdir, backbone, metrics):
model = ImageClassifier(10, backbone=backbone, metrics=metrics)
train_dl = torch.utils.data.DataLoader(DummyDataset())
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
trainer.finetune(model, train_dl, strategy="freeze_unfreeze")
def test_load_from_checkpoint_dependency_error():
with pytest.raises(ModuleNotFoundError,
match=re.escape("'lightning-flash[image]'")):
ImageClassifier.load_from_checkpoint("not_a_real_checkpoint.pt")
train_data=train_dataset.x,
train_targets=torch.from_numpy(train_dataset.y.astype(int)),
val_data=val_dataset.x,
val_targets=torch.from_numpy(val_dataset.y.astype(int)),
transform=transform,
)
model = ImageClassifier(
backbone="resnet18",
training_strategy="prototypicalnetworks",
training_strategy_kwargs={
"epoch_length": 10 * 16,
"meta_batch_size": 4,
"num_tasks": 200,
"test_num_tasks": 2000,
"ways": datamodule.num_classes,
"shots": 1,
"test_ways": 5,
"test_shots": 1,
"test_queries": 15,
},
optimizer=torch.optim.Adam,
learning_rate=0.001,
)
trainer = flash.Trainer(
max_epochs=200,
gpus=2,
accelerator="ddp_shared",
precision=16,
)
import flash
from flash.core.data.utils import download_data
from flash.image import ImageClassificationData, ImageClassifier
# 1. Create the DataModule
download_data("https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip",
"./data")
datamodule = ImageClassificationData.from_folders(
train_folder="data/hymenoptera_data/train/",
val_folder="data/hymenoptera_data/val/",
)
# 2. Build the task
model = ImageClassifier(backbone="resnet18",
num_classes=datamodule.num_classes)
# 3. Create the trainer and finetune the model
trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
trainer.finetune(model, datamodule=datamodule, strategy="freeze")
# 4. Predict what's on a few images! ants or bees?
predictions = model.predict([
"data/hymenoptera_data/val/bees/65038344_52a45d090d.jpg",
"data/hymenoptera_data/val/bees/590318879_68cf112861.jpg",
"data/hymenoptera_data/val/ants/540543309_ddbb193ee5.jpg",
])
print(predictions)
# 5. Save the model!
trainer.save_checkpoint("image_classification_model.pt")
test_dataset = fo.Dataset.from_dir(
dataset_dir="data/hymenoptera_data/test/",
dataset_type=fo.types.ImageClassificationDirectoryTree,
)
# 3 Load data into Flash
datamodule = ImageClassificationData.from_fiftyone(
train_dataset=train_dataset,
val_dataset=val_dataset,
test_dataset=test_dataset,
batch_size=4,
)
# 4 Fine tune model
model = ImageClassifier(
backbone="resnet18",
labels=datamodule.labels,
)
trainer = flash.Trainer(
max_epochs=1,
gpus=torch.cuda.device_count(),
limit_train_batches=1,
limit_val_batches=1,
)
trainer.finetune(
model,
datamodule=datamodule,
strategy=("freeze_unfreeze", 1),
)
trainer.save_checkpoint("image_classification_model.pt")
# 5 Predict from checkpoint on data with ground truth
train_files, train_targets = load_data('train')
test_files, test_targets = load_data('test')
datamodule = ImageClassificationData.from_files(
train_files=train_files,
train_targets=train_targets,
test_files=test_files,
test_targets=test_targets,
val_split=0.1, # Use 10 % of the train dataset to generate validation one.
image_size=(128, 128),
)
# 3. Build the model
model = ImageClassifier(
backbone="resnet18",
num_classes=len(genres),
multi_label=True,
metrics=F1(num_classes=len(genres)),
)
# 4. Create the trainer. Train on 2 gpus for 10 epochs.
trainer = flash.Trainer(max_epochs=10)
# 5. Train the model
trainer.finetune(model, datamodule=datamodule, strategy="freeze")
# 6. Predict what's on a few images!
# Serialize predictions as labels, low threshold to see more predictions.
model.serializer = Labels(genres, multi_label=True, threshold=0.25)
predictions = model.predict([
"data/movie_posters/predict/tt0085318.jpg",
def run(self, best_model_path: str):
# Load the model from the model path
model = ImageClassifier.load_from_checkpoint(best_model_path)
model.serve(output="labels")
# 3.a Optional: Register a custom backbone
# This is useful to create new backbone and make them accessible from `ImageClassifier`
@ImageClassifier.backbones(name="resnet18")
def fn_resnet(pretrained: bool = True):
model = torchvision.models.resnet18(pretrained)
# remove the last two layers & turn it into a Sequential model
backbone = nn.Sequential(*list(model.children())[:-2])
num_features = model.fc.in_features
# backbones need to return the num_features to build the head
return backbone, num_features
# 3.b Optional: List available backbones
print(ImageClassifier.available_backbones())
# 4. Build the model
model = ImageClassifier(backbone="resnet18",
num_classes=datamodule.num_classes,
serializer=Labels())
# 5. Create the trainer
trainer = flash.Trainer(max_epochs=1,
limit_train_batches=1,
limit_val_batches=1)
# 6. Train the model
trainer.finetune(model,
datamodule=datamodule,
strategy=FreezeUnfreeze(unfreeze_epoch=1))
def test_active_learning_training(simple_datamodule, initial_num_labels,
query_size):
seed_everything(42)
if initial_num_labels == 0:
with pytest.warns(UserWarning,
match="No labels provided for the initial step"):
active_learning_dm = ActiveLearningDataModule(
simple_datamodule,
initial_num_labels=initial_num_labels,
query_size=query_size,
val_split=0.5,
)
else:
active_learning_dm = ActiveLearningDataModule(
simple_datamodule,
initial_num_labels=initial_num_labels,
query_size=query_size,
val_split=0.5,
)
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, num_sanity_val_steps=0)
active_learning_loop = ActiveLearningLoop(label_epoch_frequency=1,
inference_iteration=3)
active_learning_loop.connect(trainer.fit_loop)
trainer.fit_loop = active_learning_loop
trainer.finetune(model,
datamodule=active_learning_dm,
strategy="no_freeze")
# Check that all metrics are logged
assert all(
any(m in log_met
for log_met in active_learning_loop.trainer.logged_metrics)
for m in ("train", "val", "test"))
# Check that the weights has changed for both module.
classifier = active_learning_loop._lightning_module.adapter.parameters()
mc_inference = active_learning_loop.inference_model.parent_module.parameters(
)
assert all(torch.equal(p1, p2) for p1, p2 in zip(classifier, mc_inference))
if initial_num_labels == 0:
assert len(active_learning_dm._dataset) == 15
else:
assert len(active_learning_dm._dataset) == 20
assert active_learning_loop.progress.total.completed == 3
labelled = active_learning_loop.state_dict(
)["state_dict"]["datamodule_state_dict"]["labelled"]
assert isinstance(labelled, np.ndarray)
# Check that we iterate over the actual pool and that shuffle is disabled.
assert len(active_learning_dm.predict_dataloader()) == math.ceil(
(~labelled).sum() / simple_datamodule.batch_size)
assert isinstance(active_learning_dm.predict_dataloader().sampler,
SequentialSampler)
if initial_num_labels == 0:
assert len(active_learning_dm.val_dataloader()) == 4
else:
# in the second scenario we have more labelled data!
assert len(active_learning_dm.val_dataloader()) == 5
def test_serve():
model = ImageClassifier(2)
# TODO: Currently only servable once a preprocess has been attached
model._preprocess = ImageClassificationPreprocess()
model.eval()
model.serve()
from flash.audio import AudioClassificationData
from flash.core.data.utils import download_data
from flash.core.finetuning import FreezeUnfreeze
from flash.image import ImageClassifier
# 1. Create the DataModule
download_data("https://pl-flash-data.s3.amazonaws.com/urban8k_images.zip", "./data")
datamodule = AudioClassificationData.from_folders(
train_folder="data/urban8k_images/train",
val_folder="data/urban8k_images/val",
spectrogram_size=(64, 64),
)
# 2. Build the model.
model = ImageClassifier(backbone="resnet18", num_classes=datamodule.num_classes)
# 3. Create the trainer and finetune the model
trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
trainer.finetune(model, datamodule=datamodule, strategy=FreezeUnfreeze(unfreeze_epoch=1))
# 4. Predict what's on few images! air_conditioner, children_playing, siren e.t.c
predictions = model.predict(
[
"data/urban8k_images/test/air_conditioner/13230-0-0-5.wav.jpg",
"data/urban8k_images/test/children_playing/9223-2-0-15.wav.jpg",
"data/urban8k_images/test/jackhammer/22883-7-10-0.wav.jpg",
]
)
print(predictions)
return os.path.join(root, f"{file_id}.jpg")
datamodule = ImageClassificationData.from_csv(
"Id",
["Action", "Romance", "Crime", "Thriller", "Adventure"],
train_file="data/movie_posters/train/metadata.csv",
train_resolver=resolver,
val_file="data/movie_posters/val/metadata.csv",
val_resolver=resolver,
transform_kwargs={"image_size": (128, 128)},
batch_size=1,
)
# 2. Build the task
model = ImageClassifier(backbone="resnet18", labels=datamodule.labels, multi_label=datamodule.multi_label)
# 3. Create the trainer and finetune the model
trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
trainer.finetune(model, datamodule=datamodule, strategy="freeze")
# 4. Predict the genre of a few movies!
datamodule = ImageClassificationData.from_files(
predict_files=[
"data/movie_posters/predict/tt0085318.jpg",
"data/movie_posters/predict/tt0089461.jpg",
"data/movie_posters/predict/tt0097179.jpg",
],
batch_size=3,
)
predictions = trainer.predict(model, datamodule=datamodule, output="labels")
test_dataset = fo.Dataset.from_dir(
dataset_dir="data/hymenoptera_data/test/",
dataset_type=fo.types.ImageClassificationDirectoryTree,
)
# 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
# 1 Download data
download_data(
"https://label-studio-testdata.s3.us-east-2.amazonaws.com/lightning-flash/data.zip"
)
# 2. Load export data
datamodule = ImageClassificationData.from_labelstudio(
export_json="data/project.json",
data_folder="data/upload/",
val_split=0.2,
)
# 3. Fine tune a model
model = ImageClassifier(
backbone="resnet18",
num_classes=datamodule.num_classes,
)
trainer = flash.Trainer(max_epochs=3)
trainer.finetune(
model,
datamodule=datamodule,
strategy=("freeze_unfreeze", 1),
)
trainer.save_checkpoint("image_classification_model.pt")
# 4. Predict from checkpoint
model = ImageClassifier.load_from_checkpoint("image_classification_model.pt")
model.output = LabelsOutput()
datamodule = ImageClassificationData.from_files(predict_files=[