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=[