# How ContextNet and UNet Modules Work in the RIFE Architecture

> Understand how ContextNet and UNet modules enhance RIFE's video frame interpolation. Learn about multi-scale context extraction and feature fusion for accurate predictions.

- Repository: [hzwer/eccv2022-rife](https://github.com/hzwer/eccv2022-rife)
- Tags: deep-dive
- Published: 2026-03-03

---

**ContextNet extracts multi-scale, flow-warped context features from input frames, while UNet fuses these features with warped images and optical flow to predict the final interpolated frame and blending mask.**

The RIFE (Real-time Intermediate Flow Estimation) architecture by `hzwer/eccv2022-rife` achieves high-quality video frame interpolation through a three-stage pipeline: coarse flow estimation, context-aware feature extraction, and multi-modal fusion. While `IFNet` handles the initial motion estimation, the **ContextNet** and **UNet** modules collaborate to align spatial features and refine the output. Understanding their interaction is essential for adapting the model to custom resolutions or modifying the fusion behavior.

## ContextNet – Hierarchical Flow-Aware Feature Extraction

**ContextNet** (and its lightweight variant **Contextnet**) generates spatially-aligned feature pyramids that enable the UNet to reason about both fine textures and global scene structure. The module processes each input frame independently, warping features at multiple scales using the intermediate optical flow predicted by `IFNet`.

### Architecture Implementation in RIFE_HDv2

The high-capacity version resides in [`model/oldmodel/RIFE_HDv2.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/oldmodel/RIFE_HDv2.py) and uses a channel factor `c = 32`:

```python
class ContextNet(nn.Module):
    def __init__(self):
        super(ContextNet, self).__init__()
        self.conv0 = Conv2(3, c)          # 3→c   (c=32)

        self.conv1 = Conv2(c, c)          # c→c

        self.conv2 = Conv2(c, 2*c)        # c→2c

        self.conv3 = Conv2(2*c, 4*c)      # 2c→4c

        self.conv4 = Conv2(4*c, 8*c)      # 4c→8c

    def forward(self, x, flow):
        x = self.conv0(x)
        x = self.conv1(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f1 = warp(x, flow)                # warp at 1/4 resolution

        x = self.conv2(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f2 = warp(x, flow)                # warp at 1/8 resolution

        x = self.conv3(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f3 = warp(x, flow)                # warp at 1/16 resolution

        x = self.conv4(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f4 = warp(x, flow)                # warp at 1/32 resolution

        return [f1, f2, f3, f4]

```

Each `Conv2` block performs stride-2 downsampling while doubling the channel depth, creating a pyramid with resolutions at 1/4, 1/8, 1/16, and 1/32 of the original.

### Lightweight Variant in refine.py

The newer implementation in [`model/refine.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/refine.py) uses the same logic with reduced capacity (`c = 16`) to improve inference speed:

```python
c = 16
class Contextnet(nn.Module):
    def __init__(self):
        super(Contextnet, self).__init__()
        self.conv1 = Conv2(3, c)
        self.conv2 = Conv2(c, 2*c)
        self.conv3 = Conv2(2*c, 4*c)
        self.conv4 = Conv2(4*c, 8*c)

    def forward(self, x, flow):
        x = self.conv1(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f1 = warp(x, flow)
        x = self.conv2(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f2 = warp(x, flow)
        x = self.conv3(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f3 = warp(x, flow)
        x = self.conv4(x)
        flow = F.interpolate(flow, scale_factor=0.5, mode='bilinear', align_corners=False) * 0.5
        f4 = warp(x, flow)
        return [f1, f2, f3, f4]

```

### Flow-Guided Warping at Multiple Scales

After every downsampling stage, **ContextNet** applies the differentiable `warp` function (implemented in [`model/warplayer.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/warplayer.py)) to align features with the intermediate frame. The flow tensor is progressively halved using `F.interpolate` and scaled by `0.5` to maintain correct motion vector magnitudes as spatial resolution decreases. This produces four warped feature maps that preserve semantic context despite large motion between frames.

## UNet – Fusion and Refinement Network

The **UNet** (class `Unet` in [`model/refine.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/refine.py), also referred to as `FusionNet` in earlier versions) aggregates information from the warped frames, predicted flow, soft masks, and the multi-scale context features generated by ContextNet.

### Encoder Path with Context Injection

The encoder consists of four downsampling blocks. The first layer receives a 17-channel tensor concatenating both original frames, their warped versions, the soft mask, and the bidirectional flow:

```python
class Unet(nn.Module):
    def __init__(self):
        super(Unet, self).__init__()
        self.down0 = Conv2(17, 2*c)      # 3+3+3+3+1+4 = 17 channels

        self.down1 = Conv2(4*c, 4*c)
        self.down2 = Conv2(8*c, 8*c)
        self.down3 = Conv2(16*c, 16*c)
        # ... decoder layers

    def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1):
        s0 = self.down0(torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1))
        s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
        s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
        s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
        # ... decoder forward pass

```

Each subsequent encoder stage concatenates its input with the corresponding context features from both frames (`c0[i]` and `c1[i]`), fusing image-level details with hierarchical semantic information.

### Decoder Path and Skip Connections

The decoder employs transposed convolutions (`deconv`) to upsample features while progressively reducing channel counts. Skip connections concatenate encoder outputs (`s0` through `s3`) with decoder features to recover spatial precision lost during downsampling:

```python
        x  = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
        x  = self.up1(torch.cat((x, s2), 1))
        x  = self.up2(torch.cat((x, s1), 1))
        x  = self.up3(torch.cat((x, s0), 1))
        x  = self.conv(x)  # nn.Conv2d(c, 3, 3, 1, 1)

        return torch.sigmoid(x)

```

A final `3×3` convolution projects the feature map to 3 RGB channels, and a sigmoid activation constrains outputs to `[0, 1]`. In the complete pipeline, this output represents a residual correction that refines the initial warped-frame estimate.

## Integration in the RIFE Pipeline

The top-level `Model` class in [`model/RIFE.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/RIFE.py) orchestrates the interaction between these modules:

```python

# Extract flow-aware context for both frames

c0 = self.contextnet(img0, flow[:, :2])   # Context for frame 0

c1 = self.contextnet(img1, flow[:, 2:4])  # Context for frame 1

# Upsample flow to full resolution

flow = F.interpolate(flow, scale_factor=2.0, mode='bilinear', align_corners=False) * 2.0

# Fuse everything to produce the interpolated frame

refine_output = self.fusionnet(img0, img1, flow, c0, c1, flow_gt)

```

Here, `self.contextnet` instantiates either `ContextNet` or `Contextnet` depending on the model variant, while `self.fusionnet` is the `Unet` class. The UNet receives **both raw images, their warped counterparts, the upsampled flow, and the four-scale context tensors** to predict the final interpolation.

## Practical Code Examples

### Running ContextNet on a Single Frame

```python
import torch
from model.refine import Contextnet  # or model.oldmodel.RIFE_HDv2.ContextNet

B, C, H, W = 1, 3, 256, 256
img = torch.randn(B, C, H, W).cuda()
flow = torch.randn(B, 2, H, W).cuda()  # Intermediate flow for this frame

ctx = Contextnet().cuda()
features = ctx(img, flow)  # List of 4 warped feature maps

for i, f in enumerate(features):
    print(f"Level {i}: {f.shape}")

# Output shapes: [1,16,128,128], [1,32,64,64], [1,64,32,32], [1,128,16,16]

```

### Feeding UNet with Context and Warped Images

```python
import torch
from model.refine import Unet, Contextnet, warp

B, C, H, W = 1, 3, 256, 256
img0 = torch.randn(B, C, H, W).cuda()
img1 = torch.randn(B, C, H, W).cuda()
flow = torch.randn(B, 4, H, W).cuda()  # Bidirectional flow [dx0, dy0, dx1, dy1]

# Extract context pyramids

ctx = Contextnet().cuda()
c0 = ctx(img0, flow[:, :2])
c1 = ctx(img1, flow[:, 2:4])

# Warp frames according to flow

warped0 = warp(img0, flow[:, :2])
warped1 = warp(img1, flow[:, 2:4])

# Initialize mask (learned during refinement)

mask = torch.zeros(B, 1, H, W).cuda()

# Run fusion network

unet = Unet().cuda()
output = unet(img0, img1, warped0, warped1, mask, flow, c0, c1)
print(output.shape)  # torch.Size([1, 3, 256, 256])

```

## Summary

- **ContextNet** generates four-scale feature pyramids (1/4 to 1/32 resolution) by progressively downsampling input frames and warping features with intermediate optical flow.
- **UNet** fuses these context features with raw frames, warped frames, and bidirectional flow through an encoder-decoder architecture with skip connections.
- Both modules rely on the `warp` utility from [`model/warplayer.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/warplayer.py) for differentiable backward warping.
- The high-capacity `ContextNet` (c=32) resides in [`model/oldmodel/RIFE_HDv2.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/oldmodel/RIFE_HDv2.py), while the lightweight `Contextnet` (c=16) and `Unet` are defined in [`model/refine.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/refine.py).
- Together, they enable the RIFE architecture to produce temporally consistent, high-fidelity intermediate frames by combining motion-aware feature alignment with deep feature fusion.

## Frequently Asked Questions

### What is the difference between ContextNet and Contextnet in the RIFE codebase?

**ContextNet** (capital N) refers to the high-capacity implementation in [`model/oldmodel/RIFE_HDv2.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/oldmodel/RIFE_HDv2.py) that uses 32 base channels. **Contextnet** (lowercase n) is the lightweight variant in [`model/refine.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/refine.py) using 16 base channels, designed for faster inference with minimal quality degradation. Both perform identical multi-scale feature extraction and flow-guided warping operations.

### Why does ContextNet warp features at every scale instead of only at full resolution?

Warping at each downsampling level ensures that context features remain spatially aligned with the intermediate frame despite large motion vectors. As resolution decreases, the flow magnitude is halved proportionally (`* 0.5`), allowing the network to capture both fine textures at high resolution and global motion patterns at lower resolutions. This hierarchical alignment prevents artifacts in scenes with complex motion.

### How many input channels does the UNet encoder receive and what do they represent?

The first encoder layer (`down0`) receives **17 channels**: 3 channels each for the two original RGB frames (`img0`, `img1`), 3 channels each for the flow-warped frames (`warped_img0`, `warped_img1`), 1 channel for the soft blending mask, and 4 channels for the bidirectional optical flow. Subsequent layers concatenate these with 2×, 4×, and 8× context features from both frames.

### Can the ContextNet and UNet modules be used independently of the full RIFE model?

Yes, both modules are standard PyTorch `nn.Module` classes that can be instantiated and executed independently, as shown in the code examples above. However, they require properly formatted optical flow tensors (from `IFNet` or another estimator) to perform meaningful warping. The `warp` function from [`model/warplayer.py`](https://github.com/hzwer/eccv2022-rife/blob/main/model/warplayer.py) must be available in the environment for either module to function correctly.