Implementation:Zai org CogVideo NLayerDiscriminator

Knowledge Sources	Zai_org_CogVideo
Domains	Video_Generation, Generative_Adversarial_Networks
Last Updated	2026-02-10 00:00 GMT

Overview

Implements the NLayerDiscriminator (PatchGAN) architecture and a weight initialization function for adversarial training of autoencoders, classifying local image patches as real or fake.

Description

This module provides two components for GAN-based training:

• weights_init() -- A weight initialization function that applies normal-distribution initialization to convolutional layers (mean=0, std=0.02) and batch normalization layers (weight mean=1, std=0.02; bias=0). Handles both standard nn.Conv2d and wrapped convolution modules with a .conv attribute.

• NLayerDiscriminator -- A fully convolutional PatchGAN discriminator following the Pix2Pix architecture. It builds a sequence of convolutional blocks that progressively downsample the input:
  • First layer: Conv2d(input_nc, ndf, 4, stride=2, padding=1) followed by LeakyReLU(0.2)
  • Middle layers: n_layers blocks of Conv2d with stride-2 downsampling, normalization (BatchNorm or ActNorm), and LeakyReLU(0.2). Filter count doubles at each layer up to a maximum of 8 * ndf.
  • Penultimate layer: Same structure but with stride=1 (no further downsampling)
  • Final layer: Conv2d(ndf*nf_mult, 1, 4, stride=1, padding=1) producing a single-channel spatial prediction map

The normalization layer is configurable: nn.BatchNorm2d by default, or ActNorm when use_actnorm=True. Bias usage in convolutional layers is automatically determined based on whether the normalization layer has affine parameters.

Usage

Used as the discriminator component in GeneralLPIPSWithDiscriminator for adversarial autoencoder training. The PatchGAN architecture evaluates overlapping local patches rather than the entire image, encouraging the generator to produce sharp, realistic textures at fine spatial scales.

Code Reference

Source Location

• Repository: Zai_org_CogVideo
• File: sat/sgm/modules/autoencoding/lpips/model/model.py

Signature

def weights_init(m)

class NLayerDiscriminator(nn.Module):
    def __init__(
        self,
        input_nc=3,
        ndf=64,
        n_layers=3,
        use_actnorm=False,
    )

    def forward(self, input) -> torch.Tensor

Import

from sat.sgm.modules.autoencoding.lpips.model.model import NLayerDiscriminator, weights_init

I/O Contract

Inputs

Name	Type	Required	Description
input	torch.Tensor	Yes	Image tensor of shape [B, input_nc, H, W]

Constructor Parameters

Name	Type	Required	Description
input_nc	int	No	Number of input channels. Default: 3
ndf	int	No	Base number of discriminator filters. Default: 64
n_layers	int	No	Number of downsampling convolutional layers. Default: 3
use_actnorm	bool	No	Use ActNorm instead of BatchNorm. Default: False

Outputs

Name	Type	Description
prediction_map	torch.Tensor	Spatial map of real/fake predictions, shape [B, 1, H', W'] where H' and W' depend on input size and n_layers

Usage Examples

from sat.sgm.modules.autoencoding.lpips.model.model import NLayerDiscriminator, weights_init

# Create PatchGAN discriminator with 3 downsampling layers
disc = NLayerDiscriminator(input_nc=3, ndf=64, n_layers=3)
disc.apply(weights_init)

# Get patch-level predictions
prediction_map = disc(image_tensor)  # [B, 1, H', W']

# Adversarial loss (e.g., hinge loss)
real_loss = -torch.mean(torch.min(prediction_map - 1, torch.zeros_like(prediction_map)))

Related Pages

• Principle:Zai_org_CogVideo_PatchGAN_Discriminator