text: import os
  
  from opt import get_opts
  import torch
  from collections import defaultdict
  
  from torch.utils.data import DataLoader
  from datasets import dataset_dict
  
  # models
  from models.nerf import *
  from models.rendering import *
  
  # optimizer, scheduler, visualization
  from utils import *
  
  # losses
  from losses import loss_dict
  
  # metrics
  from metrics import *
  
  # pytorch-lightning
  from pytorch_lightning.callbacks import ModelCheckpoint
  from pytorch_lightning import LightningModule, Trainer
  from pytorch_lightning.loggers import TestTubeLogger
  from pytorch_lightning.plugins import DDPPlugin
  
  
  class NeRFSystem(LightningModule):
      def __init__(self, hparams):
          super().__init__()
          self.save_hyperparameters(hparams)
  
          self.loss = loss_dict['nerfw'](coef=1)
  
          self.models_to_train = []
          self.embedding_xyz = PosEmbedding(hparams.N_emb_xyz-1, hparams.N_emb_xyz)
          self.embedding_dir = PosEmbedding(hparams.N_emb_dir-1, hparams.N_emb_dir)
          self.embeddings = {'xyz': self.embedding_xyz,
                             'dir': self.embedding_dir}
  
          if hparams.encode_a:
              self.embedding_a = torch.nn.Embedding(hparams.N_vocab, hparams.N_a)
              self.embeddings['a'] = self.embedding_a
              self.models_to_train += [self.embedding_a]
          if hparams.encode_t:
              self.embedding_t = torch.nn.Embedding(hparams.N_vocab, hparams.N_tau)
              self.embeddings['t'] = self.embedding_t
              self.models_to_train += [self.embedding_t]
  
          self.nerf_coarse = NeRF('coarse',
                                  in_channels_xyz=6*hparams.N_emb_xyz+3,
                                  in_channels_dir=6*hparams.N_emb_dir+3)
          self.models = {'coarse': self.nerf_coarse}
          if hparams.N_importance > 0:
              self.nerf_fine = NeRF('fine',
                                    in_channels_xyz=6*hparams.N_emb_xyz+3,
                                    in_channels_dir=6*hparams.N_emb_dir+3,
                                    encode_appearance=hparams.encode_a,
                                    in_channels_a=hparams.N_a,
                                    encode_transient=hparams.encode_t,
                                    in_channels_t=hparams.N_tau,
                                    beta_min=hparams.beta_min)
              self.models['fine'] = self.nerf_fine
          self.models_to_train += [self.models]
  
      def get_progress_bar_dict(self):
          items = super().get_progress_bar_dict()
          items.pop("v_num", None)
          return items
  
      def forward(self, rays, ts, eval=False):
          """Do batched inference on rays using chunk."""
          B = rays.shape[0]
          results = defaultdict(list)
          for i in range(0, B, self.hparams.chunk):
              rendered_ray_chunks = \
                  render_rays(self.models,
                              self.embeddings,
                              rays[i:i+self.hparams.chunk],
                              ts[i:i+self.hparams.chunk],
                              self.hparams.N_samples,
                              self.hparams.use_disp,
                              self.hparams.perturb,
                              self.hparams.noise_std,
                              self.hparams.N_importance,
                              self.hparams.chunk, # chunk size is effective in val mode
                              self.train_dataset.white_back)
  
              for k, v in rendered_ray_chunks.items():
                  results[k] += [v.cpu() if eval else v]
  
          for k, v in results.items():
              results[k] = torch.cat(v, 0)
          return results
  
      def setup(self, stage):
          dataset = dataset_dict[self.hparams.dataset_name]
          kwargs = {'root_dir': self.hparams.root_dir}
          if self.hparams.dataset_name == 'phototourism':
              kwargs['img_downscale'] = self.hparams.img_downscale
              kwargs['val_num'] = self.hparams.num_gpus
              kwargs['use_cache'] = self.hparams.use_cache
          elif self.hparams.dataset_name == 'blender':
              kwargs['img_wh'] = tuple(self.hparams.img_wh)
              kwargs['perturbation'] = self.hparams.data_perturb
          self.train_dataset = dataset(split='train', **kwargs)
          self.val_dataset = dataset(split='val', **kwargs)
  
      def configure_optimizers(self):
          self.optimizer = get_optimizer(self.hparams, self.models_to_train)
          scheduler = get_scheduler(self.hparams, self.optimizer)
          return [self.optimizer], [scheduler]
  
      def train_dataloader(self):
          return DataLoader(self.train_dataset,
                            shuffle=True,
                            num_workers=4,
                            batch_size=self.hparams.batch_size,
                            pin_memory=True)
  
      def val_dataloader(self):
          return DataLoader(self.val_dataset,
                            shuffle=False,
                            num_workers=4,
                            batch_size=1, # validate one image (H*W rays) at a time
                            pin_memory=True)
      
      def training_step(self, batch, batch_nb):
          rays, rgbs, ts = batch['rays'], batch['rgbs'], batch['ts']
          results = self(rays, ts)
          loss_d = self.loss(results, rgbs)
          loss = sum(l for l in loss_d.values())
  
          with torch.no_grad():
              typ = 'fine' if 'rgb_fine' in results else 'coarse'
              psnr_ = psnr(results[f'rgb_{typ}'], rgbs)
  
          self.log('lr', get_learning_rate(self.optimizer))
          self.log('train/loss', loss)
          for k, v in loss_d.items():
              self.log(f'train/{k}', v, prog_bar=True)
          self.log('train/psnr', psnr_, prog_bar=True)
  
          return loss
  
      def validation_step(self, batch, batch_nb):
          rays, rgbs, ts = batch['rays'], batch['rgbs'], batch['ts']
          rays = rays.squeeze() # (H*W, 3)
          rgbs = rgbs.squeeze() # (H*W, 3)
          ts = ts.squeeze() # (H*W)
          WH = batch['img_wh']
  
          # rgbs = rgbs.detach().clone()
          # del batch
          results = self(rays, ts, eval=True)
          for k in ['weights_coarse', 'opacity_coarse', 'rgb_coarse', 'transient_sigmas', 'beta', 'rgb_fine', 'depth_fine']:
              results[k] = results[k].to(self.device)
          loss_d = self.loss(results, rgbs)
          loss = sum(l for l in loss_d.values())
          log = {'val_loss': loss}
          typ = 'fine' if 'rgb_fine' in results else 'coarse'
      
          if batch_nb == 0:
              if self.hparams.dataset_name == 'phototourism':
                  W, H = WH[0, 0].item(), WH[0, 1].item()
              else:
                  W, H = self.hparams.img_wh
              img = results[f'rgb_{typ}'].view(H, W, 3).permute(2, 0, 1).cpu() # (3, H, W)
              img_gt = rgbs.view(H, W, 3).permute(2, 0, 1).cpu() # (3, H, W)
              depth = visualize_depth(results[f'depth_{typ}'].view(H, W)) # (3, H, W)
              stack = torch.stack([img_gt, img, depth]) # (3, 3, H, W)
              self.logger.experiment.add_images('val/GT_pred_depth',
                                                 stack, self.global_step)
  
          psnr_ = psnr(results[f'rgb_{typ}'], rgbs)
          log['val_psnr'] = psnr_
  
          return log
  
      def validation_epoch_end(self, outputs):
          mean_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
          mean_psnr = torch.stack([x['val_psnr'] for x in outputs]).mean()
  
          self.log('val/loss', mean_loss)
          self.log('val/psnr', mean_psnr, prog_bar=True)
  
  
  def main(hparams):
      system = NeRFSystem(hparams)
      checkpoint_callback = \
          ModelCheckpoint(dirpath=os.path.join(hparams.save_path, 'ckpts', hparams.exp_name),
                          filename='{epoch:d}',
                          monitor='val/psnr',
                          mode='max',
                          save_top_k=-1)
  
      logger = TestTubeLogger(save_dir=os.path.join(hparams.save_path, 'logs'),
                              name=hparams.exp_name,
                              debug=False,
                              create_git_tag=False,
                              log_graph=False)
  
      trainer = Trainer(max_epochs=hparams.num_epochs,
                        checkpoint_callback=True,
                        callbacks=[checkpoint_callback],
                        plugins=[DDPPlugin(find_unused_parameters=False)],
                        resume_from_checkpoint=hparams.ckpt_path,
                        logger=logger,
                        weights_summary=None,
                        progress_bar_refresh_rate=hparams.refresh_every,
                        gpus=hparams.num_gpus,
                        accelerator='ddp' if hparams.num_gpus>1 else None,
                        num_sanity_val_steps=1,
                        benchmark=True,
                        profiler="simple" if hparams.num_gpus==1 else None)
  
      trainer.fit(system)
  
  
  if __name__ == '__main__':
      hparams = get_opts()
      main(hparams)