LRM 代码精读 3-[models]

Models 是最重要的部分 计划仔细讲一下

计划分别花一天看rendering, models 还会顺便看看dino的架构

block.py

1. Basic attention block

LayerNorm1 + Self + LayerNorm2 +

MLP :

(2* (Linear(Inner_dim → mlp_ratio倍 再回归)+Dropout))

2. Cross attention block(take in cond)

同时每个attend或mlp之前有residual

LayerNorm1 + Cross + LayerNorm2 + Self + Layernorm3 +

MLP :

(2* (Linear(Inner_dim → mlp_ratio倍 再回归)+Dropout))

注意cross-attention的定义 包含kdim和vdim 都是condition的值

self.cross_attn = nn.MultiheadAttention(
   embed_dim=inner_dim, num_heads=num_heads, kdim=cond_dim, vdim=cond_dim,
	 dropout=attn_drop, bias=attn_bias, batch_first=True)

最后总结一下nn.MultiheadAttention所需参数

embed_dim, num_heads, kdim(opt), vdim(opt),

dropout, bias, batch_first()

**这里k,v 的dim其实可以不一样

atten(q,k,v,need_weights)

Return type → self_attend()[0]

3. Condition Modulation Bloack

包含condition + modulation(:调整) vectors

mod会作用在norm上

将layernorm 替换成了ModLN

class ModLN(nn.Module):
    """
    Modulation with adaLN(adaptive layer norm zero (adaLN-Zero) conditioning).
    
    References:
    DiT: https://github.com/facebookresearch/DiT/blob/main/models.py#L101
    """
    def __init__(self, inner_dim: int, mod_dim: int, eps: float):
        super().__init__()
        self.norm = nn.LayerNorm(inner_dim, eps=eps)
        self.mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(mod_dim, inner_dim * 2),
        )

    @staticmethod
    def modulate(x, shift, scale):
        # x: [N, L, D]
        # shift, scale: [N, D]
        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

    def forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:
        shift, scale = self.mlp(mod).chunk(2, dim=-1)  # [N, D]
        return self.modulate(self.norm(x), shift, scale)  # [N, L, D]

根据mod生成shift&scale

这里用modulation做norm

原本implement在DIT中

Scalable Diffusion Models with Transformers (DiT)

We train latent diffusion models, replacing the commonly-used U-Net backbone with a transformer that operates on latent patches

TransformerDecoder

process the input(with condition and modulation)

assert a, ”b” assertion error of b

from accelerate.logging import get_logger

logger = get_logger(__name__)
logging.basicConfig(level=logging.DEBUG, 
format='%(asctime)s - %(levelname)s - %(message)s')

#set 更高level 阻止一般debug的log
logging.basicConfig(level=logging.INFO)

logger.debug() 申明当前状态

logger.info()

create all blocks with the same block type and put them in forward pass one by one

create all block instances with(dim and cond_dim) with partial

    def forward_layer(self, layer: nn.Module, x: torch.Tensor, cond: torch.Tensor, mod: torch.Tensor):
        if self.block_type == 'basic':
            return layer(x)
        elif self.block_type == 'cond':
            return layer(x, cond)
        elif self.block_type == 'mod':
            return layer(x, mod)
        else:
            return layer(x, cond, mod)

    def forward(self, x: torch.Tensor, cond: torch.Tensor = None, mod: torch.Tensor = None):
        # x: [N, L, D]
        # cond: [N, L_cond, D_cond] or None
        # mod: [N, D_mod] or None
        self.assert_runtime_integrity(x, cond, mod)
        for layer in self.layers:
            x = self.forward_layer(layer, x, cond, mod)
        x = self.norm(x)
        return x

LRM model

def __init__(self, camera_embed_dim: int, rendering_samples_per_ray: int,
                 transformer_dim: int, transformer_layers: int, transformer_heads: int,
                 triplane_low_res: int, triplane_high_res: int, triplane_dim: int,
                 encoder_freeze: bool = True, encoder_type: str = 'dino',
                 encoder_model_name: str = 'facebook/dino-vitb16', encoder_feat_dim: int = 768):

1. Get encoder

self.encoder = self._encoder_fn(encoder_type)(
            model_name=encoder_model_name:'facebook/dino-vitb16',
            freeze=encoder_freeze:True,
      )
        
@staticmethod
def _encoder_fn(encoder_type: str):
    encoder_type = encoder_type.lower()
    assert encoder_type in ['dino', 'dinov2'], "Unsupported encoder type"
    if encoder_type == 'dino':
        from .encoders.dino_wrapper import DinoWrapper
        logger.info("Using DINO as the encoder")
      return DinoWrapper
    elif encoder_type == 'dinov2':
        from .encoders.dinov2_wrapper import Dinov2Wrapper
        logger.info("Using DINOv2 as the encoder")
      return Dinov2Wrapper

Camera embedding

self.camera_embedder = CameraEmbedder(
            raw_dim=12+4, embed_dim=camera_embed_dim,
        )

Convert raw camera dim(为什么12+4) to camera_embed

from the DIT without pos_embedding

class TimestepEmbedder(nn.Module):
		"""
		Embeds scalar timesteps into vector representations.
		"""
		def **init**(self, hidden_size, frequency_embedding_size=256):
			super().**init**()
			self.mlp = nn.Sequential(
			nn.Linear(frequency_embedding_size, hidden_size, bias=True),
			nn.SiLU(),
			nn.Linear(hidden_size, hidden_size, bias=True),
			)
			self.frequency_embedding_size = frequency_embedding_size
		
		@staticmethod
		def timestep_embedding(t, dim, max_period=10000):
		    """
		    Create sinusoidal timestep embeddings.
		    :param t: a 1-D Tensor of N indices, one per batch element.
		                      These may be fractional.
		    :param dim: the dimension of the output.
		    :param max_period: controls the minimum frequency of the embeddings.
		    :return: an (N, D) Tensor of positional embeddings.
		    """
		    # <https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py>
		    half = dim // 2
		    freqs = torch.exp(
		        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
		    ).to(device=t.device)
		    args = t[:, None].float() * freqs[None]
		    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
		    if dim % 2:
		        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
		    return embedding
		
		def forward(self, t):
		    t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
		    t_emb = self.mlp(t_freq)
		    return t_emb
		

Pos embedding

self.pos_embed = nn.Parameter(torch.randn

(1, 3*triplane_low_res**2, transformer_dim)

***** (1. / transformer_dim) ** 0.5)

1, total_pixiel_count * 3(planes), transformer_dim

triplane_low_res → H, W of each triplane image

camera_dim → transformer_dim

总结一下

1. image_feature = encoder(image) (dino) → encoder_feat_dim

2. camera_embedding = camera_embedder(camera)

D_cam_raw(16) → camera_embed_dim

1. Transformer

pos_embed → (cond_MHA + modulation_norm)

(N, L = triplane_dim= 3*width^2, transformer_dim)

→ (N, L, )

1. Plane generation

x = tokens.view(N, 3, H, W, transformer_dim)

回归triplane图片 dim

x = x.contiguous().view(3 * N, transformer_dim, H, W)  # Shape: [3*N, transformer_dim, H, W]

Upsampling

*self*.upsampler = nn.ConvTranspose2d(transformer_dim, triplane_dim, kernel_size=2, stride=2, padding=0)

triplane_dim = 3维 每张图片 深度dim (例如64)

kernal size & stride = 2 Double image size

x = self.upsampler(x)  # [3*N, D', H', W']
x = x.view(3, N, *x.shape[-3:])  # [3, N, D', H', W']
x = torch.einsum('indhw->nidhw', x)  # [N, 3, D', H', W']

planes → (N, 3, D, H*, W*) H*, W* = 2H, 2W

1. Synthesize(Render)

render_results = *self*.synthesizer(planes, render_cameras, render_anchors, render_resolutions, render_bg_colors, render_region_size)

render_results = self.synthesizer(
planes,
render_cameras,       # Shape: [N, M, D_cam_render]  #rendering cam parameter(specific viewpoints)
render_anchors,       # Shape: [N, M, 2]
render_resolutions,   # Shape: [N, M, 1]
render_bg_colors,     # Shape: [N, M, 1]
render_region_size    # Integer
)

Output → rendered images (N, M(camera pos count), C_img(3?), H, W)

Summary of Dimensions at Each Stage

1. Input Image:
   • Shape: [N, 3, H_img, W_img]
2. Image Features:
   • After Encoder: [N, encoder_feat_dim]
3. Camera Embeddings:
   • After Camera Embedder: [N, camera_embed_dim]
4. Transformer Input:
   • Positional Embeddings Repeated: [N, L, transformer_dim]
5. Transformer Output (Tokens):
   • Shape: [N, L, transformer_dim]
6. Tokens Reshaped for Planes:
   • Before Upsampling: [N, 3, H, W, transformer_dim]
   • After Reordering: [3*N, transformer_dim, H, W]
7. Upsampled Planes:
   • After Upsampler: [3*N, triplane_dim, 2*H, 2*W]
8. Final Planes:
   • Reshaped Back: [N, 3, triplane_dim, 2*H, 2*W]
9. Rendered Images:
   • Shape: [N, M, C_img, H_render, W_render]

Role of Camera and Image Dimensions

• Image Dimensions:
  • The input images are encoded into a compact feature representation, capturing the visual content.
  • The feature vector is used to condition the transformer, influencing the generation of triplanar features.
• Camera Dimensions:
  • Raw camera parameters are embedded into a feature space.
  • These embeddings modulate the transformer, allowing it to account for viewpoint information during plane generation.
• Transformer Integration:
  • The transformer combines positional embeddings, image features, and camera embeddings to produce tokens that represent triplanar features.
  • The sequence length L is directly related to the resolution of the triplanes.
• Plane Generation:
  • Tokens are reshaped and upsampled to create high-resolution triplanar feature planes.
  • These planes encapsulate 3D information and are used for rendering novel views.

https://www.bilibili.com/video/BV1B6421F7iH?spm_id_from=333.788.recommend_more_video.-1&vd_source=143a2ef0cd4b513f15da9430a5ab01fc