Global module variable re-initialized at training time

[mfoglio]

I have a complex data structure that requires adding every sample of my dataset to a dictionary _annotation_ids_str_to_int.
This variable is common to the training, validation, and test splits. It is defined as a global variable in a module.
The variable seems to be initialized correctly. However, at training time, the variable needs to be accessed by the __getitem__ function. This function finds the variable to be empty.
I am not sure if the variable has been reinitialized or the issue is caused by multithreading or multiprocessing.
How do I get the __getitem__ function to properly access _annotation_ids_str_to_int ?
Here's a dummy code that you can use to reproduce the error. It is composed of two files. Note the print calls that will help you to find the issue.
Also, note that the variable _annotation_ids_str_to_int is not exactly a dictionary, but it is instead a bidictionary which is simply a dictionary that offers a mapping from value to key when accessing my_dict.inverse.
main.py

from my_datamodule import SampleDataModule
import torch
import pytorch_lightning as pl

import utils


log = utils.get_pylogger(__name__)


class Net(torch.nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.layer = torch.nn.Linear(10, 2)

    def forward(self, x):
        return self.layer(x)


class Model(pl.LightningModule):

    def __init__(self, net):
        super(Model, self).__init__()
        self.net = net
        self.loss = torch.nn.MSELoss()
        self.optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
        self.scheduler = torch.optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=0.9)

    def forward(self, x):
        return self.net(x)

    def training_step(self, batch, batch_id):
        inputs, y = batch
        y_pred = self.forward(inputs)
        loss = self.loss(y_pred, y)
        return {'loss': loss}

    def validation_step(self, batch, batch_id):
        inputs, y = batch
        y_pred = self.forward(inputs)
        loss = self.loss(y_pred, y)
        return {'loss': loss}

    def configure_optimizers(self):
        pass


def main():
    net = Net()
    model = Model(net)
    datamodule = SampleDataModule()
    trainer = pl.Trainer()
    trainer.fit(model, datamodule=datamodule)


if __name__ == "__main__":
    main()

my_datamodule.py

from typing import Dict, Optional

from bidict import bidict
import torch
from torch.utils.data import DataLoader, Dataset
from pytorch_lightning import LightningDataModule


_annotation_ids_str_to_int: Dict[str, int] = bidict()


class BaseImageDataset(Dataset):

    def __init__(self):
        self.length = torch.randint(5, 10, (1, )).item()
        for i in range(self.length):
            _annotation_ids_str_to_int[i] = (torch.tensor([i] * 10), i * 2)
        print("__init__ Current keys are", _annotation_ids_str_to_int.keys())

    def __len__(self):
        return self.length

    def __getitem__(self, index: int):
        print("__getitem__ Current keys are", _annotation_ids_str_to_int.keys())
        return _annotation_ids_str_to_int[index]


class SampleDataModule(LightningDataModule):

    def __init__(self):
        super().__init__()
        self.save_hyperparameters(logger=False)
        self.train_dataset = None
        self.val_dataset = None
        self.test_dataset = None

    def prepare_data(self):
        pass

    def setup(self, stage: Optional[str] = None):
        # Initialize dataset splits
        self.train_dataset = BaseImageDataset()
        self.val_dataset = BaseImageDataset()
        self.test_dataset = BaseImageDataset()

    @property
    def _get_data_loaders_common_kwargs(self):
        return dict(
            batch_size=8,
            num_workers=1,
            pin_memory=True,
            persistent_workers=True,
            prefetch_factor=1,
            drop_last=False
        )

    def train_dataloader(self):
        return DataLoader(self.train_dataset, **self._get_data_loaders_common_kwargs)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, **self._get_data_loaders_common_kwargs)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, **self._get_data_loaders_common_kwargs)

The output is the following:

GPU available: False, used: False
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/loops/utilities.py:89: PossibleUserWarning: `max_epochs` was not set. Setting it to 1000 epochs. To train without an epoch limit, set `max_epochs=-1`.
  rank_zero_warn(
Missing logger folder: /Users/myuser/project//src/issue_torch_lightning/lightning_logs
/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py:182: UserWarning: `LightningModule.configure_optimizers` returned `None`, this fit will run with no optimizer
  rank_zero_warn(
  | Name | Type    | Params
---------------------------------
0 | net  | Net     | 22    
1 | loss | MSELoss | 0     
---------------------------------
22        Trainable params
0         Non-trainable params
22        Total params
0.000     Total estimated model params size (MB)
**__init__ Current keys are dict_keys([0, 1, 2, 3, 4, 5, 6, 7])
__init__ Current keys are dict_keys([0, 1, 2, 3, 4, 5, 6, 7])
__init__ Current keys are dict_keys([0, 1, 2, 3, 4, 5, 6, 7])**
Sanity Checking: 0it [00:00, ?it/s]/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:236: PossibleUserWarning: The dataloader, val_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(
**__getitem__ Current keys are dict_keys([])**
Traceback (most recent call last):
  File "/Users/matteo/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/223.8836.43/PyCharm.app/Contents/plugins/python/helpers/pydev/pydevconsole.py", line 364, in runcode
    coro = func()
  File "<input>", line 1, in <module>
  File "/Users/matteo/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/223.8836.43/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 198, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "/Users/matteo/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/223.8836.43/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/Users/myuser/project//src/issue_torch_lightning/test_bug_torch_lightning.py", line 58, in <module>
    main()
  File "/Users/myuser/project//src/issue_torch_lightning/test_bug_torch_lightning.py", line 54, in main
    trainer.fit(model, datamodule=datamodule)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 696, in fit
    self._call_and_handle_interrupt(
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 650, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 735, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1166, in _run
    results = self._run_stage()
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1252, in _run_stage
    return self._run_train()
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1274, in _run_train
    self._run_sanity_check()
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1343, in _run_sanity_check
    val_loop.run()
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 200, in run
    self.advance(*args, **kwargs)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 155, in advance
    dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 200, in run
    self.advance(*args, **kwargs)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 127, in advance
    batch = next(data_fetcher)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/utilities/fetching.py", line 184, in __next__
    return self.fetching_function()
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/utilities/fetching.py", line 263, in fetching_function
    self._fetch_next_batch(self.dataloader_iter)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/pytorch_lightning/utilities/fetching.py", line 277, in _fetch_next_batch
    batch = next(iterator)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 652, in __next__
    data = self._next_data()
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1347, in _next_data
    return self._process_data(data)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1373, in _process_data
    data.reraise()
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/torch/_utils.py", line 461, in reraise
    raise exception
KeyError: Caught KeyError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 302, in _worker_loop
    data = fetcher.fetch(index)
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 49, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/Users/myuser/project//src/issue_torch_lightning/my_datamodule.py", line 25, in __getitem__
    return _annotation_ids_str_to_int[index]
  File "/Users/myuser/project//venv/lib/python3.8/site-packages/bidict/_base.py", line 523, in __getitem__
    return self._fwdm[key]
KeyError: 0

I am using python 3.8.14 on macOS with the following requirements:

bidict==0.22.1
pytorch-lightning==1.7.7
pytorch-lightning-bolts==0.3.2.post1
sagemaker-pytorch-training==2.6.2.post0
torch==1.12.0
torchmetrics==0.10.0
torchvision==0.13.0

[mfoglio]

I solved this by adding torch.multiprocessing.set_start_method('fork'). I think this is due to the way multiprocess works in python: see https://pytorch.org/docs/stable/data.html#platform-specific-behaviors .