PoseFormer
Pipeline
2D skeleton sequence --> Spatial Transformer --> Temporal Transformer --> Regression Head.
2D skeleton sequence
\(X\in R^{f\times (2J)}\), \(f\): 接受帧数, 感受野; \(J\): 关节点数; \(2\): 2D空间\((x, y)\).
\(x^i \in R^{1\times (2J)}\), 单个骨架关节点坐标向量.
Spatial Transformer
The spatial transformer module is to extract a high dimension feature embedding
from a single frame.
Spatial emedding(joint embedding)
对于单帧2D姿态\(x\in R^{1\times (2J)}\), Spatial Transformer Module把每个2D坐标\((x, y)\)作为
一个patch(ViT), 用一个线性投影将2D坐标投影至高纬空间. 接着加上表示关节位置信息的embedding,
作为输入交给Spatial Transformer学习关节点间空间信息.
\(x_i \in R^{1\times (2J)}\) -- linear transformation --> \(x_i \in R^{J\times c}\) -- positional embeddomh --> \(z_0^i \in R^{J\times C}\)
\(z_0 = [p^1E; p^2E; ...; p^JE] + E_{Spos}\), 其中\(p\)表示单个2D坐标\((x, y)\),
\(E\)表示线性投影矩阵, \(E_{Spos}\)表示positional embedding.
Self-attention
将\(z_0\in R^{f\times c}\)作为输入进入Spatial Transformer模块, 由\(L\)个ViT模块堆叠而成.
\(z'_l = MSA(LN(z_{l-1})) + Z_{l - 1}, l = 1, 2, ..., L\)
\(z_l = MLP(LN(Z'_l) + Z'_l), l = 1, 2, ..., L\)
\(Y = LN(Z_L)\)
其中\(LN()\)表示layer normalization, \(Y\in R^(f\times c)\), 与输入维度相同.
即先对输入进行多头注意力, 再进入全连接. 每个输入进入网络前先进行层归一化, 输出网络后残差连接.
Temporal Transformer
The goal of the temporal transformer module is to model dependencies across the
sequence of frames.
对于第\(i\)帧2D骨架经过Spatial Transformer后输出为\(z_L^i \in R^{J\times c}\), 将其
压缩为向量\(z^i\in R^{1\times (Jc)}\). 将\(f\)个(感受野)向量连接concatenate, 加上表示
帧位置信息的position embedding作为输入
\(Z_0 = [z_L^1; z_L^2; ..., z_L^f] + E_{Tpos}, Z_0\in R^{f\times (Jc)}\).
将\(Z_0\)输入\(L\)层相同的ViT模块(与Spatial Transformer相同), 得到输出\(Y\in R^{f\times (Jc)}\).
Regression Head
由于最终预测中间一帧的3D骨架, 首先需要将输出维度缩减为\(f\times (Jc) --> 1\times (Jc)\).
本方法通过对帧维度(\(f\))进行可训练的带权平均weighted mean操作实现. 最后, 使用一个简单的
\(MLP\), 得到最终输出\(y\in R^{1\times (3J)}\), 代表3D骨架坐标向量.
Code
Spatial Transformer与Temporal Transformer均以ViT作为blocks串联作为模块.
Spatial Transformer的输入为\(J\times c\), \(J\)表示关节点个数, \(c\)表示Spatial embedding的维度.
(关节点作为token)此时\(B\times F\)作为batch size, 其中\(B\)是mini batch数, \(F\)是感受野, 即2D骨架序列长度.
Spatial Transformer的输入为\(F\times Jc\), 将骨架信息作为token, \(B\)作为Batch size.
import math
import logging
from functools import partial
from collections import OrderedDict
from einops import rearrange, repeat
import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.models.helpers import load_pretrained
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from timm.models.registry import register_model
class Mlp(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class Attention(nn.Module):
def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
# NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
self.scale = qk_scale or head_dim ** -0.5
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
# qkv: [3, B, num_heads, N, C // num_heads]
q, k, v = qkv[0], qkv[1], qkv[2] # [B, num_heads, N, C // num_heads] make torchscript happy (cannot use tensor as tuple)
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x) # linear, C -- > C
x = self.proj_drop(x)
return x
# ViT
class Block(nn.Module):
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
# dim 单个输入向量的长度 N个向量组成矩阵 B个batch
super().__init__()
self.norm1 = norm_layer(dim)
self.attn = Attention(
dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
def forward(self, x):
x = x + self.drop_path(self.attn(self.norm1(x)))
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x
class PoseTransformer(nn.Module):
def __init__(self, num_frame=9, num_joints=17, in_chans=2, out_dim=17*3, spatial_embed_dim=32, depth=4,
num_heads=8, mlp_ratio=2., qkv_bias=True, qk_scale=None,
drop_rate=0., attn_drop_rate=0., drop_path_rate=0.2, norm_layer=None):
""" ##########hybrid_backbone=None, representation_size=None,
Args:
num_frame (int, tuple): input frame number
num_joints (int, tuple): joints number
in_chans (int): number of input channels, 2D joints have 2 channels: (x,y)
embed_dim_ratio (int): embedding dimension ratio
depth (int): depth of transformer
num_heads (int): number of attention heads
mlp_ratio (int): ratio of mlp hidden dim to embedding dim
qkv_bias (bool): enable bias for qkv if True
qk_scale (float): override default qk scale of head_dim ** -0.5 if set
drop_rate (float): dropout rate
attn_drop_rate (float): attention dropout rate
drop_path_rate (float): stochastic depth rate
norm_layer: (nn.Module): normalization layer
"""
super().__init__()
self.out_dim = out_dim
norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
temporal_embed_dim = spatial_embed_dim * num_joints #### temporal embed_dim is num_joints * spatial embedding dim ratio
#out_dim = num_joints * 3 #### output dimension is num_joints * 3
### spatial patch embedding (x, y) --> embedding
self.Spatial_patch_to_embedding = nn.Linear(in_chans, spatial_embed_dim)
# 1 x J x embedding. add as joint position info
self.Spatial_pos_embed = nn.Parameter(torch.zeros(1, num_joints, spatial_embed_dim))
self.Temporal_pos_embed = nn.Parameter(torch.zeros(1, num_frame, temporal_embed_dim))
self.pos_drop = nn.Dropout(p=drop_rate)
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
self.Spatial_blocks = nn.ModuleList([
Block(
dim=spatial_embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
for i in range(depth)])
self.Temporal_blocks = nn.ModuleList([
Block(
dim=temporal_embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
for i in range(depth)])
self.Spatial_norm = norm_layer(spatial_embed_dim)
self.Temporal_norm = norm_layer(temporal_embed_dim)
####### A easy way to implement weighted mean
self.weighted_mean = torch.nn.Conv1d(in_channels=num_frame, out_channels=1, kernel_size=1)
self.head = nn.Sequential(
nn.LayerNorm(temporal_embed_dim),
nn.Linear(temporal_embed_dim , out_dim),
)
def Spatial_forward_features(self, x):
# input shape: B x 2 x F x J
b, _, f, j = x.shape
# B x 2 x F x J --> BF x J x 2
x = rearrange(x, 'b c f j -> (b f) j c', )
# BF x J x 2 --> BF x J x c, 2 -- Linear Transformation --> c
x = self.Spatial_patch_to_embedding(x)
x += self.Spatial_pos_embed # pos_embed: a matrix (J x c)
x = self.pos_drop(x)
# BF x J x c --> BF x J x c, J x c as input to ViT. c as token, BF as batch size
for blk in self.Spatial_blocks:
x = blk(x)
x = self.Spatial_norm(x)
# BF x J x c --> B x F x Jc
x = rearrange(x, '(b f) j c -> b f (j c)', f=f)
return x
def Temporal_forward_features(self, x):
b = x.shape[0]
x += self.Temporal_pos_embed # pos_embed: matrix (F x Jc)
x = self.pos_drop(x)
# B x F x Jc --> B x F x Jc, F x Jc as input, Js as token, B as batch size
for blk in self.Temporal_blocks:
x = blk(x)
x = self.Temporal_norm(x)
# B x F x Jc --> B x 1 x Jc. average of F frames
x = self.weighted_mean(x)
# B x 1 x Jc --> B x 1 x Jc
x = x.view(b, 1, -1)
return x
def forward(self, x):
# B x F x J x 2 --> B x 2 x F x J
x = x.permute(0, 3, 1, 2)
b, _, _, _ = x.shape
j = self.out_dim // 3
# B x 2 x F x J --> B x F x Jc, c is spatial embedding dimension
x = self.Spatial_forward_features(x)
# B x F x Jc --> B x 1 x Jc
x = self.Temporal_forward_features(x)
# B x 1 x Jc --> B x 1 x 3J
x = self.head(x) # Jc --> 3J. spatial embedding --> 3 D joints
# B x 1 x 3J --> B x 1 x J x 3
x = x.view(b, 1, j, -1) # batch size, 1, j, 3 (1, j, 3): center frame joints
return x
if __name__ == '__main__':
model = PoseTransformer(num_frame=27)
input = torch.randn(27, 27, 17, 2)
output = model(input)
print(output.shape)
