This is an independent PyTorch implementation of SRT, as presented in the paper "Geometry-Free Novel View Synthesis Through Set-Latent Scene Representations" by Sajjadi et al., CVPR 2022. All credit for the model goes to the original authors.

After cloning the repository and creating a new conda environment, the following steps will get you started:

The code currently supports the following datasets. Simply download and place (or symlink) them in the data directory.

• The 3D datasets introduced by ObSuRF.

• The NMR multiclass ShapeNet dataset, hosted by Niemeyer et al. It may be downloaded here.

• SRT's MultiShapeNet (MSN) dataset, specifically version 2.3.3. It may be downloaded via gsutil:

pip install gsutil
mkdir -p data/msn/multi_shapenet_frames/
gsutil -m cp -r gs://kubric-public/tfds/multi_shapenet_frames/2.3.3/ data/msn/multi_shapenet_frames/

This code requires at least Python 3.9 and PyTorch 1.11. Additional dependencies may be installed via pip -r requirements.txt. Note that Tensorflow is required to load SRT's MultiShapeNet data, though the CPU version suffices.

Rendering videos additionally depends on ffmpeg>=4.3 being available in your $PATH.

Each run's config, checkpoints, and visualization are stored in a dedicated directory. Recommended configs can be found under runs/[dataset]/[model].

To train a model on a single GPU, simply run e.g.:

python train.py runs/nmr/srt/config.yaml

To train on multiple GPUs on a single machine, launch multiple processes via Torchrun, where $NUM_GPUS is the number of GPUs to use:

torchrun --standalone --nnodes 1 --nproc_per_node $NUM_GPUS train.py runs/nmr/srt/config.yaml

Checkpoints are automatically stored in and (if available) loaded from the run directory. Visualizations and evaluations are produced periodically. Check the args of train.py for additional options. Importantly, to log training progress, use the --wandb flag to enable Weights & Biases.

The default training configurations require a significant amount of total VRAM.

• SRT on NMR (nmr/srt/config.yaml) requires about 130GB, e.g. 4 A100 GPUs with 40GB VRAM, each.
• SRT on MSN (msn/srt/config.yaml) requires about 350GB, e.g. 6 A100 GPUS with 80GB VRAM, each.

If you do not have those resources, consider modifying the config files by reducing the batch size (training: batch_size), the number of target points per scene (data: num_points), or both. The model has not appeared to be particularly sensitive to either.

Videos may be rendered using render.py, e.g.

python render.py runs/nmr/srt/config.yaml --sceneid 1 --motion rotate_and_closeup --fade

Rendered frames and videos are placed in the run directory. Check the args of render.py for various camera movements, and compile_video.py for different ways of compiling videos.

Here you find some checkpoints which partially reproduce the quantitative results of the paper.

(*) The MSN run was largely trained with a batch size of 192, due to memory constrains.

If you have the compute to spare to train SRT for the full 4 million iterations, please reach out and let me know how it goes!

On the NMR dataset, SRT overfits to the 24 camera positions in the training dataset (left). It will generally not produce coherent images when given other cameras at test time (right).

On the MSN dataset, this is not an issue, as cameras are densely sampled in the training data. The model even has some ability to produce closeups which use cameras outside of the training distribution.

@article{srt22,
  title={{Scene Representation Transformer: Geometry-Free Novel View Synthesis Through Set-Latent Scene Representations}},
  author={Mehdi S. M. Sajjadi and Henning Meyer and Etienne Pot and Urs Bergmann and Klaus Greff and Noha Radwan and Suhani Vora and Mario Lucic and Daniel Duckworth and Alexey Dosovitskiy and Jakob Uszkoreit and Thomas Funkhouser and Andrea Tagliasacchi},
  journal={{CVPR}},
  year={2022},
  url={https://srt-paper.github.io/},
}