A question on the discriminator detail about stylegan-v HOT 6 CLOSED

from typing import Tuple
import torch
import torch.nn as nn
import numpy as np
from omegaconf import DictConfig


class TemporalDifferenceEncoder(nn.Module):
    def __init__(self, cfg: DictConfig):
        super().__init__()

        self.cfg = cfg

        if self.cfg.sampling.num_frames_per_video > 1:
            self.d = 256
            self.const_embed = nn.Embedding(self.cfg.sampling.max_num_frames, self.d)
            self.time_encoder = ScalarEncoder1d(self.cfg.sampling.max_num_frames)

    def get_dim(self) -> int:
        if self.cfg.sampling.num_frames_per_video == 1:
            return 1
        elif self.cfg.sampling.type == 'uniform':
            return self.d + self.time_encoder.get_dim()
        else:
            return (self.d + self.time_encoder.get_dim()) * (self.cfg.sampling.num_frames_per_video - 1)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        batch_size = t.shape[0]

        if self.cfg.sampling.num_frames_per_video == 1:
            out = torch.zeros(len(t), 1, device=t.device)
        else:
            if self.cfg.sampling.type == 'uniform':
                num_diffs_to_use = 1
                t_diffs = t[:, 1] - t[:, 0] # [batch_size]
            else:
                num_diffs_to_use = self.cfg.sampling.num_frames_per_video - 1
                t_diffs = (t[:, 1:] - t[:, :-1]).view(-1) # [batch_size * (num_frames - 1)]
            # Note: float => round => long is necessary when it's originally long
            const_embs = self.const_embed(t_diffs.float().round().long()) # [batch_size * num_diffs_to_use, d]
            fourier_embs = self.time_encoder(t_diffs.unsqueeze(1)) # [batch_size * num_diffs_to_use, num_fourier_feats]
            out = torch.cat([const_embs, fourier_embs], dim=1) # [batch_size * num_diffs_to_use, d + num_fourier_feats]
            out = out.view(batch_size, num_diffs_to_use, -1).view(batch_size, -1) # [batch_size, num_diffs_to_use * (d + num_fourier_feats)]

        return out

#----------------------------------------------------------------------------

class ScalarEncoder1d(nn.Module):
    def __init__(self, max_num_frames: int):
        super().__init__()

        assert max_num_frames >= 1, f"Wrong max_num_frames: {max_num_frames}"
        fourier_coefs = construct_log_spaced_freqs(max_num_frames)
        self.register_buffer('fourier_coefs', fourier_coefs) # [1, num_fourier_feats]

    def get_dim(self) -> int:
        return self.fourier_coefs.shape[1] * 2

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        assert t.ndim == 2, f"Wrong shape: {t.shape}"

        t = t.view(-1).float() # [batch_size * num_frames]
        fourier_raw_embs = self.fourier_coefs * t.unsqueeze(1) # [bf, num_fourier_feats]

        fourier_embs = torch.cat([
            fourier_raw_embs.sin(),
            fourier_raw_embs.cos(),
        ], dim=1) # [bf, num_fourier_feats * 2]

        return fourier_embs

#----------------------------------------------------------------------------

def construct_log_spaced_freqs(max_num_frames: int) -> Tuple[int, torch.Tensor]:
    time_resolution = 2 ** np.ceil(np.log2(max_num_frames))
    num_fourier_feats = np.ceil(np.log2(time_resolution)).astype(int)
    powers = torch.tensor([2]).repeat(num_fourier_feats).pow(torch.arange(num_fourier_feats)) # [num_fourier_feats]
    fourier_coefs = powers.unsqueeze(0).float() * np.pi # [1, num_fourier_feats]

    return fourier_coefs / time_resolution

#----------------------------------------------------------------------------

class PositionalEncoding(Module):
    """
    TODO: Other PE type implementation, E.g. [..., cos(2*pi*sigma^(j/m)*v), sin(2*pi*sigma^(j/m)*v), ...]
    """
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.omega = nn.Parameter(torch.randn(1, self.dim, requires_grad=True))
        self.rho   = nn.Parameter(torch.randn(self.dim, requires_grad=True))
        self.alpha = nn.Parameter(torch.randn(self.dim, requires_grad=True))
    
    def forward(self, Ts):
        """
        Ts: [b, t, 1]
        """
        ot = torch.einsum("btc,cd->btd", [Ts, self.omega])
        pe = torch.sin(ot + self.rho) * self.alpha
        return pe