Encountered 'NoneType' object has no attribute 'canvas' error while training a semantic segmentation model

[gtgrp-user]

I encountered the following error while training a semantic segmentation model using a GeoDataModule.

Constants used for the training:

accelerator = "gpu" if torch.cuda.is_available() else "cpu"
default_root_dir = os.path.join('./', "full_experiments_geodataset_v1")
checkpoint_callback = ModelCheckpoint(
    monitor="val_loss", dirpath=default_root_dir, save_top_k=1, save_last=True
)
early_stopping_callback = EarlyStopping(monitor="val_loss", min_delta=0.00, patience=10)
logger = TensorBoardLogger(save_dir=default_root_dir, name="tutorial_logs")

num_workers = 4
max_epochs = 100
fast_dev_run = False

torch.set_float32_matmul_precision('high')

Creating the semantic segmentation task and appropriate trainer

task = SemanticSegmentationTask(
    loss="jaccard",
    model="deeplabv3+",
    backbone="resnet50",
    weights='imagenet',
    in_channels=10,
    num_classes=7,
    learning_rate=0.001,
    learning_rate_schedule_patience=5,
    ignore_index=None
)

trainer = Trainer(
    accelerator=accelerator,
    callbacks=[checkpoint_callback, early_stopping_callback],
    enable_progress_bar = True,
    fast_dev_run=fast_dev_run,
    log_every_n_steps=1,
    logger=logger,
    min_epochs=20,
    max_epochs=max_epochs,
)

Error encountered in this method

trainer.fit(task, datamodule = datamodule)

Stack trace:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[21], line 1
----> 1 trainer.fit(task, datamodule = datamodule)

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py:608, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    606 model = self._maybe_unwrap_optimized(model)
    607 self.strategy._lightning_module = model
--> 608 call._call_and_handle_interrupt(
    609     self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    610 )

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/call.py:38, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     36         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
     37     else:
---> 38         return trainer_fn(*args, **kwargs)
     40 except _TunerExitException:
     41     trainer._call_teardown_hook()

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py:650, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    643 ckpt_path = ckpt_path or self.resume_from_checkpoint
    644 self._ckpt_path = self._checkpoint_connector._set_ckpt_path(
    645     self.state.fn,
    646     ckpt_path,  # type: ignore[arg-type]
    647     model_provided=True,
    648     model_connected=self.lightning_module is not None,
    649 )
--> 650 self._run(model, ckpt_path=self.ckpt_path)
    652 assert self.state.stopped
    653 self.training = False

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py:1112, in Trainer._run(self, model, ckpt_path)
   1108 self._checkpoint_connector.restore_training_state()
   1110 self._checkpoint_connector.resume_end()
-> 1112 results = self._run_stage()
   1114 log.detail(f"{self.__class__.__name__}: trainer tearing down")
   1115 self._teardown()

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py:1191, in Trainer._run_stage(self)
   1189 if self.predicting:
   1190     return self._run_predict()
-> 1191 self._run_train()

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py:1204, in Trainer._run_train(self)
   1201 self._pre_training_routine()
   1203 with isolate_rng():
-> 1204     self._run_sanity_check()
   1206 # enable train mode
   1207 assert self.model is not None

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py:1276, in Trainer._run_sanity_check(self)
   1274 # run eval step
   1275 with torch.no_grad():
-> 1276     val_loop.run()
   1278 self._call_callback_hooks("on_sanity_check_end")
   1280 # reset logger connector

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/loops/loop.py:199, in Loop.run(self, *args, **kwargs)
    197 try:
    198     self.on_advance_start(*args, **kwargs)
--> 199     self.advance(*args, **kwargs)
    200     self.on_advance_end()
    201     self._restarting = False

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/loops/dataloader/evaluation_loop.py:152, in EvaluationLoop.advance(self, *args, **kwargs)
    150 if self.num_dataloaders > 1:
    151     kwargs["dataloader_idx"] = dataloader_idx
--> 152 dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
    154 # store batch level output per dataloader
    155 self._outputs.append(dl_outputs)

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/loops/loop.py:199, in Loop.run(self, *args, **kwargs)
    197 try:
    198     self.on_advance_start(*args, **kwargs)
--> 199     self.advance(*args, **kwargs)
    200     self.on_advance_end()
    201     self._restarting = False

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/loops/epoch/evaluation_epoch_loop.py:137, in EvaluationEpochLoop.advance(self, data_fetcher, dl_max_batches, kwargs)
    134 self.batch_progress.increment_started()
    136 # lightning module methods
--> 137 output = self._evaluation_step(**kwargs)
    138 output = self._evaluation_step_end(output)
    140 self.batch_progress.increment_processed()

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/loops/epoch/evaluation_epoch_loop.py:234, in EvaluationEpochLoop._evaluation_step(self, **kwargs)
    223 """The evaluation step (validation_step or test_step depending on the trainer's state).
    224 
    225 Args:
   (...)
    231     the outputs of the step
    232 """
    233 hook_name = "test_step" if self.trainer.testing else "validation_step"
--> 234 output = self.trainer._call_strategy_hook(hook_name, *kwargs.values())
    236 return output

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py:1494, in Trainer._call_strategy_hook(self, hook_name, *args, **kwargs)
   1491     return
   1493 with self.profiler.profile(f"[Strategy]{self.strategy.__class__.__name__}.{hook_name}"):
-> 1494     output = fn(*args, **kwargs)
   1496 # restore current_fx when nested context
   1497 pl_module._current_fx_name = prev_fx_name

File /usr/local/lib/python3.10/dist-packages/lightning/pytorch/strategies/strategy.py:390, in Strategy.validation_step(self, *args, **kwargs)
    388 with self.precision_plugin.val_step_context():
    389     assert isinstance(self.model, ValidationStep)
--> 390     return self.model.validation_step(*args, **kwargs)

File /usr/local/lib/python3.10/dist-packages/torchgeo/trainers/segmentation.py:211, in SemanticSegmentationTask.validation_step(self, *args, **kwargs)
    209     fig = datamodule.plot(sample)
    210     summary_writer = self.logger.experiment
--> 211     summary_writer.add_figure(
    212         f"image/{batch_idx}", fig, global_step=self.global_step
    213     )
    214     plt.close()
    215 except ValueError:

File /usr/local/lib/python3.10/dist-packages/torch/utils/tensorboard/writer.py:750, in SummaryWriter.add_figure(self, tag, figure, global_step, close, walltime)
    740     self.add_image(
    741         tag,
    742         figure_to_image(figure, close),
   (...)
    745         dataformats="NCHW",
    746     )
    747 else:
    748     self.add_image(
    749         tag,
--> 750         figure_to_image(figure, close),
    751         global_step,
    752         walltime,
    753         dataformats="CHW",
    754     )

File /usr/local/lib/python3.10/dist-packages/torch/utils/tensorboard/_utils.py:35, in figure_to_image(figures, close)
     33     return np.stack(images)
     34 else:
---> 35     image = render_to_rgb(figures)
     36     return image

File /usr/local/lib/python3.10/dist-packages/torch/utils/tensorboard/_utils.py:24, in figure_to_image.<locals>.render_to_rgb(figure)
     22 canvas.draw()
     23 data = np.frombuffer(canvas.buffer_rgba(), dtype=np.uint8)
---> 24 w, h = figure.canvas.get_width_height()
     25 image_hwc = data.reshape([h, w, 4])[:, :, 0:3]
     26 image_chw = np.moveaxis(image_hwc, source=2, destination=0)

AttributeError: 'NoneType' object has no attribute 'canvas'

[adamjstewart]

Which datamodule? Does the dataset for this datamodule have a plot function? Need more details to reproduce this bug.

[gtgrp-user]

I have used a custom datamodule derived from GeoDataModule and copied the setup method from L7Irish. Used RasterDataset for this custom datamodule. I don't have a plot function for the dataset, I am afraid.

class REGeoDataModule(GeoDataModule):
    # Copied from L7Irish
    def setup(self, stage: str) -> None:
        """Set up datasets.

        Args:
            stage: Either 'fit', 'validate', 'test', or 'predict'.
        """
        self.dataset = self.dataset_class(**self.kwargs)
        
        generator = torch.Generator().manual_seed(0)
        (
            self.train_dataset,
            self.val_dataset,
            self.test_dataset,
        ) = random_bbox_assignment(self.dataset, [0.6, 0.2, 0.2], generator)
        
        if stage in ["fit"]:
            self.train_batch_sampler = RandomBatchGeoSampler(
                self.train_dataset, self.patch_size, self.batch_size, self.length
            )
        if stage in ["fit", "validate"]:
            self.val_sampler = GridGeoSampler(
                self.val_dataset, self.patch_size, self.patch_size
            )
        if stage in ["test"]:
            self.test_sampler = GridGeoSampler(
                self.test_dataset, self.patch_size, self.patch_size
            )

[adamjstewart]

Thanks, I think this bug was introduced in #992 when I switched from catching AttributeError to ValueError. Our datamodules check to see whether or not a plot method is defined, but our trainers don't check to see whether or not a real figure was returned. Let me see if I can fix this.

[adamjstewart]

Hmm, all of the tests pass when I try a dataset without a plot function. So I'm still not sure how to reproduce this error. It's easy to fix, but if I can't test it, then I can't guarantee it won't break again. What versions of matplotlib/tensorboard/etc. are you using? You can usually get this from pip list.

[adamjstewart]

Do you override __init__ in your datamodule? If not, how do you instantiate your datamodule? Can you share your dataset implementation? Still struggling to reproduce this issue.

[gtgrp-user]

I am using matplotlib v3.7.2 and tensorboard v2.13.0, I have also attached the full list here requirements.txt

[adamjstewart]

If you comment out the plot function in SEN12MS and run the segmentation trainer tests, do they fail? That's the test I'm trying locally. Basically:

1. Clone the repo
2. Comment out plot in torchgeo/datasets/sen12ms.py
3. Run pytest tests/trainers/test_segmentation.py

[gtgrp-user]

One test is marked as skipped, while the remaining tests appear to be successful.

Running 42 items in this shard

tests/trainers/test_segmentation.py ..........s...............................                                                                                                                    [100%]

============================================================================= 41 passed, 1 skipped, 23 deselected in 12.85s =============================================================================

[adamjstewart]

Okay that's good to know. I'll keep trying to figure out a way to reproduce this bug in our tests so I can guarantee that it won't come back after I fix it.

[gtgrp-user]

Thanks, much appreciated.

[adamjstewart]

By the way, in the meantime, you can just add a plot function to your dataset that returns an empty figure. Not sure if that was clear or not.

[adamjstewart]

Can you open an issue for this so I don't forget it? I'm still looking for ways to reproduce the issue in our tests...

[gtgrp-user]

Do you mean the GeoDatamodule without a PLOT issue?

[adamjstewart]

Actually, it looks like there's a "create issue from discussion" button, let me just use that.

[adamjstewart]

Created #1551