model.predict from example code does not work

🐛 Bug

model.predict using the sample code on my own data seems to fail. Not really sure how to interpret this error message.

KeyError: None
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-42-c24524f3eb71> in <module>
     14 # ])
     15 
---> 16 print(model.predict("classifier/test/bad/Image__2021-07-26__12-52-29_aligned-18.jpg"))

~\miniconda3\envs\torch\lib\site-packages\flash\core\model.py in wrapper(self, *args, **kwargs)
     76         torch.set_grad_enabled(False)
     77 
---> 78         result = func(self, *args, **kwargs)
     79 
     80         if is_training:

~\miniconda3\envs\torch\lib\site-packages\flash\core\model.py in predict(self, x, data_source, deserializer, data_pipeline)
    234             x = self.transfer_batch_to_device(x, self.device, 0)
    235         else:
--> 236             x = self.transfer_batch_to_device(x, self.device)
    237         x = data_pipeline.device_preprocessor(running_stage)(x)
    238         predictions = self.predict_step(x, 0)  # batch_idx is always 0 when running with `model.predict`

~\miniconda3\envs\torch\lib\site-packages\flash\core\data\data_pipeline.py in __call__(self, *args, **kwargs)
    599             stage = self.model.trainer.state.stage
    600 
--> 601         internal_running_state = self.internal_mapping[stage]
    602         additional_func = self._stage_mapping.get(internal_running_state, None)
    603 

KeyError: None

To Reproduce

Steps to reproduce the behavior:

1. Follow instructions here
2. Observe error

Code sample

from flash.core.classification import Probabilities

model.serializer = Probabilities()

print(model.predict("classifier/some image here..."))

Expected behavior

Example expected to work. model.predict seems to be the culprit. I can train and evaluate just fine.

Environment

• PyTorch Version (e.g., 1.0): 1.9
• OS (e.g., Linux): Windows 10
• How you installed PyTorch (conda, pip, source): conda
• Build command you used (if compiling from source):
• Python version: 3.8

[ethanwharris]

Awesome, thanks for that! I’ll look into it and get back to you once I have something 😃

[dlangerm]

@ethanwharris Here’s the snippet that triggered the bug

import flash
from flash.image import ImageClassificationData, ImageClassifier
from torchmetrics import Accuracy, F1, Precision, Recall, Specificity, MetricCollection
from pytorch_lightning import seed_everything
from flash.core.classification import Labels

seed_everything(42)

# 2. Load the data, everything within these folders is separated into /good /bad 
datamodule = ImageClassificationData.from_folders(
    train_folder="classifier/train/",
    val_folder="classifier/val/",
    test_folder="classifier/test/",
    # removing this triggers the bug
    # predict_folder="classifier/test/",
)

# 3. Build the model
model = ImageClassifier(
    num_classes=datamodule.num_classes, 
    backbone="resnet18", 
    metrics=MetricCollection([
        Accuracy(num_classes=2),
        F1(num_classes=2, average='macro'),
        Recall(num_classes=2),
        Precision(num_classes=2),
        Specificity(num_classes=2, average='macro'),
    ]),
)

# 4. Create the trainer. Run once on data
trainer = flash.Trainer(
    max_epochs=25,
    gpus=1,
)

# 5. Finetune the model
trainer.finetune(model, datamodule=datamodule, strategy="freeze")
model.serializer = Labels()
print(model.predict(["classifier/test/bad/XXX some image here"]))

lightning version 1.3.8, flash version 0.4.0