PyTorch Lightning script for logging with Weights&Biases in offline mode and automatic sync at the end of training

pytorch_lightning_wandb_offline_sync.py

import pytorch_lightning as pl
import wandb
from os.path import dirname
import subprocess

from pytorch_lightning.loggers.wandb import WandbLogger
from pytorch_lightning.utilities.cli import LightningCLI


class WandBOfflineSync(pl.LightningModule):
  def __init__(self, offline: bool = True):
    super().__init__()
    self.save_hyperparameters()
  def on_fit_end(self) -> None:
    if self.hparams.offline is True:
        # Syncing W&B at the end
        # 1. save sync dir (after marking a run finished, the W&B object changes (is teared down?)
        sync_dir = dirname(self.logger.experiment.dir)
        # 2. mark run complete
        wandb.finish()
        # 3. call the sync command for the run directory
        subprocess.check_call(["wandb", "sync", sync_dir])
        
        
  class MyLightningCLI(LightningCLI):
    def before_fit(self):
        if isinstance(self.trainer.logger, WandbLogger) is True:
            if self.config[self.subcommand].model.offline is True:
                self.trainer.logger.__dict__["_wandb_init"]["mode"] = "offline"
            else:
                self.trainer.logger.__dict__["_wandb_init"]["mode"] = "online"