| import torch
|
| import torch.nn as nn
|
|
|
|
|
| class FactorConv3d(nn.Module):
|
| """
|
| (2+1)D 分解 3D 卷积：1×H×W 空间卷积 → Swish → T×1×1 时间卷积
|
| """
|
| def __init__(self,
|
| in_channels: int,
|
| out_channels: int,
|
| kernel_size,
|
| stride: int = 1,
|
| dilation: int = 1):
|
| super().__init__()
|
|
|
| if isinstance(kernel_size, int):
|
| k_t, k_h, k_w = kernel_size, kernel_size, kernel_size
|
| else:
|
| k_t, k_h, k_w = kernel_size
|
|
|
| pad_t = (k_t - 1) * dilation // 2
|
| pad_hw = (k_h - 1) * dilation // 2
|
|
|
| self.spatial = nn.Conv3d(
|
| in_channels, in_channels,
|
| kernel_size=(1, k_h, k_w),
|
| stride=(1, stride, stride),
|
| padding=(0, pad_hw, pad_hw),
|
| dilation=(1, dilation, dilation),
|
| groups=in_channels,
|
| bias=False
|
| )
|
|
|
| self.temporal = nn.Conv3d(
|
| in_channels, out_channels,
|
| kernel_size=(k_t, 1, 1),
|
| stride=(stride, 1, 1),
|
| padding=(pad_t, 0, 0),
|
| dilation=(dilation, 1, 1),
|
| bias=True
|
| )
|
|
|
| self.act = nn.SiLU()
|
|
|
| def forward(self, x):
|
| x = self.spatial(x)
|
| x = self.act(x)
|
| x = self.temporal(x)
|
| return x
|
|
|
|
|
| class LayerNorm2D(nn.Module):
|
| """
|
| LayerNorm over C for a 4-D tensor (B, C, H, W)
|
| """
|
| def __init__(self, num_channels, eps=1e-5, affine=True):
|
| super().__init__()
|
| self.num_channels = num_channels
|
| self.eps = eps
|
| self.affine = affine
|
| if affine:
|
| self.weight = nn.Parameter(torch.ones(1, num_channels, 1, 1))
|
| self.bias = nn.Parameter(torch.zeros(1, num_channels, 1, 1))
|
|
|
| def forward(self, x):
|
|
|
| mean = x.mean(dim=1, keepdim=True)
|
| var = x.var (dim=1, keepdim=True, unbiased=False)
|
| x = (x - mean) / torch.sqrt(var + self.eps)
|
| if self.affine:
|
| x = x * self.weight + self.bias
|
| return x
|
|
|
|
|
| class PoseRefNetNoBNV3(nn.Module):
|
| def __init__(self,
|
| in_channels_c: int,
|
| in_channels_x: int,
|
| hidden_dim: int = 256,
|
| num_heads: int = 8,
|
| dropout: float = 0.1):
|
| super().__init__()
|
| self.d_model = hidden_dim
|
| self.nhead = num_heads
|
|
|
| self.proj_p = nn.Conv2d(in_channels_c, hidden_dim, kernel_size=1)
|
| self.proj_r = nn.Conv2d(in_channels_x, hidden_dim, kernel_size=1)
|
|
|
| self.proj_p_back = nn.Conv2d(hidden_dim, in_channels_c, kernel_size=1)
|
|
|
| self.cross_attn = nn.MultiheadAttention(hidden_dim,
|
| num_heads=num_heads,
|
| dropout=dropout)
|
|
|
| self.ffn_pose = nn.Sequential(
|
| nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1),
|
| nn.SiLU(),
|
| nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1)
|
| )
|
|
|
| self.norm1 = LayerNorm2D(hidden_dim)
|
| self.norm2 = LayerNorm2D(hidden_dim)
|
|
|
| def forward(self, pose, ref, mask=None):
|
| """
|
| pose : (B, C1, T, H, W)
|
| ref : (B, C2, T, H, W)
|
| mask : (B, T*H*W) 可选 key_padding_mask
|
| return: (B, d_model, T, H, W)
|
| """
|
| B, _, T, H, W = pose.shape
|
| L = H * W
|
|
|
| p_trans = pose.permute(0, 2, 1, 3, 4).contiguous().flatten(0, 1)
|
| r_trans = ref.permute(0, 2, 1, 3, 4).contiguous().flatten(0, 1)
|
|
|
| p_trans = self.proj_p(p_trans)
|
| r_trans = self.proj_r(r_trans)
|
|
|
| p_trans = p_trans.flatten(2).transpose(1, 2)
|
| r_trans = r_trans.flatten(2).transpose(1, 2)
|
|
|
| out = self.cross_attn(query=r_trans,
|
| key=p_trans,
|
| value=p_trans,
|
| key_padding_mask=mask)[0]
|
|
|
| out = out.transpose(1, 2).contiguous().view(B*T, -1, H, W)
|
| out = self.norm1(out)
|
|
|
| ffn_out = self.ffn_pose(out)
|
| out = out + ffn_out
|
| out = self.norm2(out)
|
| out = self.proj_p_back(out)
|
| out = out.view(B, T, -1, H, W).contiguous().transpose(1, 2)
|
|
|
| return out
|
|
|
