# How to Train RIFE on a Custom Video Dataset Using Distributed Multi-GPU Training

> Train RIFE on custom video data across multiple GPUs. Learn to implement a PyTorch Dataset, adapt train.py, and launch distributed training for faster model generation.

- Repository: [hzwer/eccv2022-rife](https://github.com/hzwer/eccv2022-rife)
- Tags: how-to-guide
- Published: 2026-03-03

---

**To train RIFE on custom video data with distributed multi-GPU setup, implement a PyTorch Dataset class that yields concatenated frame triplets of shape `[9, H, W]` and timestep tensors, replace `VimeoDataset` in [`train.py`](https://github.com/hzwer/eccv2022-rife/blob/main/train.py), and launch with `torch.distributed.launch` specifying `--nproc_per_node` to match your GPU count.**

The hzwer/eccv2022-rife repository provides a complete training pipeline for the RIFE (Real-Time Intermediate Flow Estimation) architecture, but adapting it to your own video collection requires modifying the data loading logic. The default implementation uses `VimeoDataset` defined in [`dataset.py`](https://github.com/hzwer/eccv2022-rife/blob/main/dataset.py), yet the distributed training framework in [`train.py`](https://github.com/hzwer/eccv2022-rife/blob/main/train.py) is designed to work with any PyTorch Dataset that returns the correct tensor shapes. By implementing a custom dataset and leveraging the existing NCCL process group initialization, you can scale training across multiple GPUs while processing your own frame sequences.

## Understanding the Distributed Training Pipeline in train.py

The distributed training orchestration resides entirely in [`train.py`](https://github.com/hzwer/eccv2022-rife/blob/main/train.py), which sets up a PyTorch DistributedDataParallel (DDP) environment using NCCL backend. At line 45, the script initializes the process group with `torch.distributed.init_process_group(backend="nccl", world_size=args.world_size)`, establishing communication across all available GPUs. Immediately following at line 46, each process binds to a specific CUDA device via `torch.cuda.set_device(args.local_rank)` to prevent device contention.

The data parallelism relies on `DistributedSampler` (line 48) to partition the dataset across processes, ensuring each GPU processes mutually exclusive subsets during an epoch. The DataLoader at line 50 wraps this sampler with `DataLoader(dataset, batch_size=args.batch_size, sampler=sampler, ...)` to load mini-batches. Synchronization occurs at line 95 where `dist.barrier()` ensures all processes complete the current epoch before proceeding to validation or checkpointing, preventing race conditions during model state saves.

## Implementing a Custom Dataset for Video Triplets

To replace `VimeoDataset`, create a class inheriting from `torch.utils.data.Dataset` that returns two specific items: a frame tensor of shape `[9, H, W]` and a timestep tensor of shape `(1, 1, 1)`. The frame tensor must concatenate the first frame `I0`, the third frame `I1`, and the ground-truth middle frame `It` along the channel dimension (channels 0-2 for `I0`, 3-5 for `I1`, 6-8 for `It`).

The timestep tensor represents the temporal position of the intermediate frame between `I0` and `I1`, typically `0.5` for halfway interpolation, but customizable for arbitrary-time interpolation tasks. Your `__getitem__` method should load three consecutive frames from your custom folder structure, apply consistent spatial cropping to ensure matching dimensions, convert from HWC to CHW format using `permute(2, 0, 1)`, and concatenate them before returning.

```python
import os
import cv2
import torch
from torch.utils.data import Dataset

class MyVideoDataset(Dataset):
    """Custom dataset yielding triplets (I0, It, I1) for RIFE training."""
    def __init__(self, root, split='train', crop_h=224, crop_w=224):
        self.root = root
        self.split = split
        self.crop_h = crop_h
        self.crop_w = crop_w
        self.seqs = [d for d in os.listdir(os.path.join(root, split)) 
                     if os.path.isdir(os.path.join(root, split, d))]

    def __len__(self):
        return len(self.seqs)

    def __getitem__(self, idx):
        seq_path = os.path.join(self.root, self.split, self.seqs[idx])
        
        # Load frames - adapt naming convention to your data

        I0 = cv2.imread(os.path.join(seq_path, 'frame_000.png'))
        It = cv2.imread(os.path.join(seq_path, 'frame_001.png'))
        I1 = cv2.imread(os.path.join(seq_path, 'frame_002.png'))
        
        # Random crop augmentation

        h, w, _ = I0.shape
        x = torch.randint(0, h - self.crop_h + 1, (1,)).item()
        y = torch.randint(0, w - self.crop_w + 1, (1,)).item()
        
        I0 = I0[x:x+self.crop_h, y:y+self.crop_w]
        It = It[x:x+self.crop_h, y:y+self.crop_w]
        I1 = I1[x:x+self.crop_h, y:y+self.crop_w]
        
        # Convert to CHW tensors

        I0 = torch.from_numpy(I0).permute(2, 0, 1).float()
        It = torch.from_numpy(It).permute(2, 0, 1).float()
        I1 = torch.from_numpy(I1).permute(2, 0, 1).float()
        
        # Concatenate as [I0, I1, It] -> shape (9, H, W)

        frames = torch.cat((I0, I1, It), dim=0)
        
        # Fixed timestep 0.5 (modify for variable interpolation)

        t = torch.tensor([0.5]).view(1, 1, 1)
        
        return frames, t

```

## Integrating Your Dataset into the Training Script

After creating your dataset class, modify [`train.py`](https://github.com/hzwer/eccv2022-rife/blob/main/train.py) to import and instantiate it instead of `VimeoDataset`. Comment out or remove the import statement `from dataset import VimeoDataset` and add `from my_dataset import MyVideoDataset`. Update the dataset instantiation around lines 48-52 to use your class while preserving the `DistributedSampler` wrapper, which remains essential for correct data sharding across GPUs.

```python

# Replace in train.py

from my_dataset import MyVideoDataset  # Your custom implementation

# Training set with distributed sampling

dataset = MyVideoDataset(root='/data/myvideos', split='train')
sampler = DistributedSampler(dataset)
train_data = DataLoader(dataset,
                        batch_size=args.batch_size,
                        num_workers=8,
                        pin_memory=True,
                        drop_last=True,
                        sampler=sampler)

# Validation set (no distributed sampler needed for validation)

dataset_val = MyVideoDataset(root='/data/myvideos', split='val')
val_data = DataLoader(dataset_val, batch_size=16, num_workers=8, pin_memory=True)

```

Maintain the `args.local_rank` parameter handling (lines 38-44 in [`train.py`](https://github.com/hzwer/eccv2022-rife/blob/main/train.py)) as the distributed launcher automatically assigns this value to each process to identify specific GPUs within the node.

## Launching Multi-GPU Distributed Training

Execute distributed training using the legacy `torch.distributed.launch` module or the modern `torchrun` utility (PyTorch ≥ 1.9). Both methods spawn independent processes that execute [`train.py`](https://github.com/hzwer/eccv2022-rife/blob/main/train.py) with automatically set `local_rank` arguments.

**Using torch.distributed.launch:**

```bash
python3 -m torch.distributed.launch --nproc_per_node=4 train.py \
    --world_size=4 \
    --batch_size=8 \
    --epoch=200

```

**Using torchrun (recommended for PyTorch 1.9+):**

```bash
torchrun --standalone --nnodes=1 --nproc_per_node=4 train.py \
    --world_size=4 \
    --batch_size=8 \
    --epoch=200

```

The `--nproc_per_node` flag must match `--world_size` for single-node training, specifying the total number of GPUs to utilize. The script automatically handles GPU assignment via `args.local_rank`, which the launcher populates for each process. Ensure your custom dataset path in [`train.py`](https://github.com/hzwer/eccv2022-rife/blob/main/train.py) points to the correct root directory containing your video frames before launching.

## Summary

- **Custom Dataset Requirements**: Implement `torch.utils.data.Dataset` returning tensors of shape `[9, H, W]` (concatenated `I0`, `I1`, `It`) and timestep `(1, 1, 1)`.
- **Key File Modifications**: Replace `VimeoDataset` imports and instantiations in [`train.py`](https://github.com/hzwer/eccv2022-rife/blob/main/train.py) with your custom class while keeping the `DistributedSampler` wrapper.
- **Distributed Setup**: The training script initializes NCCL process groups at line 45, assigns GPUs via `local_rank` at line 46, and synchronizes with `dist.barrier()` at line 95.
- **Launch Commands**: Use `torch.distributed.launch --nproc_per_node=N` or `torchrun --nproc_per_node=N` where N matches your GPU count and `--world_size` parameter.

## Frequently Asked Questions

### What exact tensor shapes does RIFE expect from the dataset?

The model requires two tensors: a frame tensor of shape `[9, H, W]` containing the first frame (channels 0-2), third frame (channels 3-5), and ground-truth middle frame (channels 6-8) concatenated along the channel dimension, and a timestep tensor of shape `(1, 1, 1)` representing the interpolation position between 0 and 1. These shapes are hardcoded in the training loop within [`train.py`](https://github.com/hzwer/eccv2022-rife/blob/main/train.py) and the RIFE model implementation in [`model/RIFE.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/RIFE.py).

### Can I train with variable timesteps instead of fixed 0.5?

Yes. Modify your Dataset's `__getitem__` method to compute and return the appropriate timestep value based on the temporal distance between your input frames. The model supports arbitrary-time interpolation when provided with variable `t` values, though the default `VimeoDataset` and standard training loops use `0.5` for midpoint interpolation. Ensure your timestep tensor maintains shape `(1, 1, 1)` regardless of the value.

### How do I handle videos with different resolutions in the same dataset?

Implement consistent spatial preprocessing in your Dataset class. The reference `VimeoDataset` in [`dataset.py`](https://github.com/hzwer/eccv2022-rife/blob/main/dataset.py) applies random cropping to fixed dimensions (lines 73-104), which you should replicate by specifying `crop_h` and `crop_w` parameters. Alternatively, implement resizing or padding logic to standardize dimensions before converting to tensors, as the model requires fixed-size inputs within a batch.

### Is single-GPU training supported without distributed launch?

Yes, though [`train.py`](https://github.com/hzwer/eccv2022-rife/blob/main/train.py) is optimized for distributed training. For single-GPU execution, you can run `python train.py --world_size=1` directly without the distributed launcher, but you must modify the code to bypass `DistributedSampler` (which requires an initialized process group) or initialize a dummy process group. For simplicity, use the distributed launcher with `--nproc_per_node=1` to maintain code compatibility.