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

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, 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, yet the distributed training framework in 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, 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.

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 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.

# 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) 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 with automatically set local_rank arguments.

Using torch.distributed.launch:

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+):

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 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 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 and the RIFE model implementation in 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 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 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}