first commit
This commit is contained in:
IamZLT
176
model/SGraFormer.py
Normal file
176
model/SGraFormer.py
Normal file
@ -0,0 +1,176 @@
|
||||
## Our model was revised from https://github.com/zczcwh/PoseFormer/blob/main/common/model_poseformer.py
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from functools import partial
|
||||
from einops import rearrange
|
||||
from timm.models.layers import DropPath
|
||||
|
||||
from common.opt import opts
|
||||
|
||||
from model.Spatial_encoder import First_view_Spatial_features, Spatial_features
|
||||
from model.Temporal_encoder import Temporal__features
|
||||
|
||||
opt = opts().parse()
|
||||
|
||||
|
||||
#######################################################################################################################
|
||||
class sgraformer(nn.Module):
|
||||
def __init__(self, num_frame=9, num_joints=17, in_chans=2, embed_dim_ratio=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__()
|
||||
|
||||
embed_dim = embed_dim_ratio * num_joints
|
||||
out_dim = num_joints * 3 #### output dimension is num_joints * 3
|
||||
##Spatial_features
|
||||
self.SF1 = First_view_Spatial_features(num_frame, num_joints, in_chans, embed_dim_ratio, depth,
|
||||
num_heads, mlp_ratio, qkv_bias, qk_scale,
|
||||
drop_rate, attn_drop_rate, drop_path_rate, norm_layer)
|
||||
self.SF2 = Spatial_features(num_frame, num_joints, in_chans, embed_dim_ratio, depth,
|
||||
num_heads, mlp_ratio, qkv_bias, qk_scale,
|
||||
drop_rate, attn_drop_rate, drop_path_rate, norm_layer)
|
||||
self.SF3 = Spatial_features(num_frame, num_joints, in_chans, embed_dim_ratio, depth,
|
||||
num_heads, mlp_ratio, qkv_bias, qk_scale,
|
||||
drop_rate, attn_drop_rate, drop_path_rate, norm_layer)
|
||||
self.SF4 = Spatial_features(num_frame, num_joints, in_chans, embed_dim_ratio, depth,
|
||||
num_heads, mlp_ratio, qkv_bias, qk_scale,
|
||||
drop_rate, attn_drop_rate, drop_path_rate, norm_layer)
|
||||
|
||||
## MVF
|
||||
self.view_pos_embed = nn.Parameter(torch.zeros(1, 4, num_frame, embed_dim))
|
||||
self.pos_drop = nn.Dropout(p=0.)
|
||||
|
||||
self.conv = nn.Sequential(
|
||||
nn.BatchNorm2d(4, momentum=0.1),
|
||||
nn.Conv2d(4, 1, kernel_size=opt.mvf_kernel, stride=1, padding=int(opt.mvf_kernel // 2), bias=False),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
|
||||
self.conv_hop = nn.Sequential(
|
||||
nn.BatchNorm2d(4, momentum=0.1),
|
||||
nn.Conv2d(4, 1, kernel_size=opt.mvf_kernel, stride=1, padding=int(opt.mvf_kernel // 2), bias=False),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
self.conv_norm = nn.LayerNorm(embed_dim)
|
||||
|
||||
self.conv_hop_norm = nn.LayerNorm(embed_dim)
|
||||
|
||||
|
||||
# Time Serial
|
||||
self.TF = Temporal__features(num_frame, num_joints, in_chans, embed_dim_ratio, depth,
|
||||
num_heads, mlp_ratio, qkv_bias, qk_scale,
|
||||
drop_rate, attn_drop_rate, drop_path_rate, norm_layer)
|
||||
|
||||
self.head = nn.Sequential(
|
||||
nn.LayerNorm(embed_dim),
|
||||
nn.Linear(embed_dim, out_dim),
|
||||
)
|
||||
|
||||
self.hop_w0 = nn.Parameter(torch.ones(17, 17))
|
||||
self.hop_w1 = nn.Parameter(torch.ones(17, 17))
|
||||
self.hop_w2 = nn.Parameter(torch.ones(17, 17))
|
||||
self.hop_w3 = nn.Parameter(torch.ones(17, 17))
|
||||
self.hop_w4 = nn.Parameter(torch.ones(17, 17))
|
||||
|
||||
self.hop_global = nn.Parameter(torch.ones(17, 17))
|
||||
|
||||
self.linear_hop = nn.Linear(8, 2)
|
||||
# self.max_pool = nn.MaxPool1d(2)
|
||||
|
||||
self.edge_embedding = nn.Linear(17*17*4, 17*17)
|
||||
|
||||
def forward(self, x, hops):
|
||||
b, f, v, j, c = x.shape
|
||||
|
||||
edge_embedding = self.edge_embedding(hops[0].reshape(1, -1))
|
||||
|
||||
###############golbal feature#################
|
||||
x_hop_global = x.unsqueeze(3).repeat(1, 1, 1, 17, 1, 1)
|
||||
x_hop_global = x_hop_global - x_hop_global.permute(0, 1, 2, 4, 3, 5)
|
||||
x_hop_global = torch.sum(x_hop_global ** 2, dim=-1)
|
||||
hop_global = x_hop_global / torch.sum(x_hop_global, dim=-1).unsqueeze(-1)
|
||||
hops = hops.unsqueeze(1).unsqueeze(2).repeat(1, f, v, 1, 1, 1)
|
||||
hops1 = hop_global * hops[:, :, :, 0]
|
||||
hops2 = hop_global * hops[:, :, :, 1]
|
||||
hops3 = hop_global * hops[:, :, :, 2]
|
||||
hops4 = hop_global * hops[:, :, :, 3]
|
||||
# hops = torch.cat((hops1,hops2,hops3,hops4), dim=-1)
|
||||
hops = torch.cat((hops1,hops2,hops3,hops4), dim=-1)
|
||||
|
||||
|
||||
x1 = x[:, :, 0]
|
||||
x2 = x[:, :, 1]
|
||||
x3 = x[:, :, 2]
|
||||
x4 = x[:, :, 3]
|
||||
|
||||
x1 = x1.permute(0, 3, 1, 2)
|
||||
x2 = x2.permute(0, 3, 1, 2)
|
||||
x3 = x3.permute(0, 3, 1, 2)
|
||||
x4 = x4.permute(0, 3, 1, 2)
|
||||
|
||||
hop1 = hops[:, :, 0]
|
||||
hop2 = hops[:, :, 1]
|
||||
hop3 = hops[:, :, 2]
|
||||
hop4 = hops[:, :, 3]
|
||||
|
||||
hop1 = hop1.permute(0, 3, 1, 2)
|
||||
hop2 = hop2.permute(0, 3, 1, 2)
|
||||
hop3 = hop3.permute(0, 3, 1, 2)
|
||||
hop4 = hop4.permute(0, 3, 1, 2)
|
||||
|
||||
### Semantic graph transformer encoder
|
||||
x1, hop1, MSA1, MSA2, MSA3, MSA4 = self.SF1(x1, hop1, edge_embedding)
|
||||
x2, hop2, MSA1, MSA2, MSA3, MSA4 = self.SF2(x2, hop2, MSA1, MSA2, MSA3, MSA4, edge_embedding)
|
||||
x3, hop3, MSA1, MSA2, MSA3, MSA4 = self.SF3(x3, hop3, MSA1, MSA2, MSA3, MSA4, edge_embedding)
|
||||
x4, hop4, MSA1, MSA2, MSA3, MSA4 = self.SF4(x4, hop4, MSA1, MSA2, MSA3, MSA4, edge_embedding)
|
||||
|
||||
### Multi-view cross-channel fusion
|
||||
x = torch.cat((x1.unsqueeze(1), x2.unsqueeze(1), x3.unsqueeze(1), x4.unsqueeze(1)), dim=1) + self.view_pos_embed
|
||||
x = self.pos_drop(x)
|
||||
x = self.conv(x).squeeze(1) + x1 + x2 + x3 + x4
|
||||
x = self.conv_norm(x)
|
||||
|
||||
hop = torch.cat((hop1.unsqueeze(1), hop2.unsqueeze(1), hop3.unsqueeze(1), hop4.unsqueeze(1)), dim=1) + self.view_pos_embed
|
||||
hop = self.pos_drop(hop)
|
||||
# hop = self.conv_hop(hop).squeeze(1) + hop1 + hop2 + hop3 + hop4
|
||||
# hop = self.conv_hop_norm(hop)
|
||||
hop = self.conv(hop).squeeze(1) + hop1 + hop2 + hop3 + hop4
|
||||
hop = self.conv_norm(hop)
|
||||
|
||||
x = x * hop
|
||||
|
||||
|
||||
### Temporal transformer encoder
|
||||
x = self.TF(x)
|
||||
|
||||
x = self.head(x)
|
||||
x = x.view(b, opt.frames, j, -1)
|
||||
|
||||
print("=============> x.shape", x.shape)
|
||||
return x
|
||||
|
||||
|
||||
# x = torch.rand((8, 27, 4, 17 , 2))
|
||||
# hops = torch.rand((8,4,17,17))
|
||||
# mvft = hmvformer(num_frame=opt.frames, num_joints=17, in_chans=2, embed_dim_ratio=32, depth=4,
|
||||
# num_heads=8, mlp_ratio=2., qkv_bias=True, qk_scale=None, drop_path_rate=0.1)
|
||||
# print(mvft(x, hops).shape)
|
||||
Reference in New Issue
Block a user