multiple processes running all tasks after trainer.fit(accelerator="ddp")
❓ Questions and Help
What is your question?
When training with ddp the script calls multiple python scripts to run the training. This causes an issue when I use the same python script to do other stuff after Im done with training.
What is best practice here? My only solution so far would be to condition on os.environ["LOCAL_RANK"]=='0' when doing stuff later on. This feels a bit hacky - is there a better way to do this?
Code
Running the example code with 2 gpus:

import os
import torch
from torch import nn
import torch.nn.functional as F
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
import pytorch_lightning as pl

class LitAutoEncoder(pl.LightningModule):

    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))

    def forward(self, x):
        # in lightning, forward defines the prediction/inference actions
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        # training_step defined the train loop. It is independent of forward
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('train_loss', loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
train, val = random_split(dataset, [55000, 5000])

autoencoder = LitAutoEncoder()
trainer = pl.Trainer(
    max_epochs=1, 
    overfit_batches = 0.05,
    gpus =-1, 
    accelerator="ddp")
trainer.fit(autoencoder, DataLoader(train), DataLoader(val))

print(f"This will print twice {os.environ['LOCAL_RANK']}")

if os.environ["LOCAL_RANK"] =='0':
    print(f"This will print once {os.environ['LOCAL_RANK']}")


What's your environment?

ubuntu
pytorch 1.6.0
pytorch-lightning 1.1.6