近期在看Shift-GCN的论文[1]，该网络是基于Shift卷积算子[2]在图结构数据上的延伸。在阅读源代码[3]的时候发现了其对于Non-Local Spatial Shift Graph Convolution有意思的实现方法，在这里简要记录一下。 本文转载自徐飞翔的"Shift-GCN中Shift的实现细节笔记，通过torch.index_select实现"
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
在讨论代码本身之前，简要介绍下Non-Local Spatial Shift Graph Convolution的操作流程，具体介绍可见博文[1]。对于一个时空骨骼点序列而言，如Fig 1所示，将单帧的骨骼点图视为是完全图，因此任何一个节点都和其他所有节点有所连接，其shift卷积策略为：
对于一个特征图
而言，其中
是骨骼点数量，
是特征通道数。对于第
个通道的
距离为
。
根据这种简单的策略，如Fig 1所示，形成了类似于螺旋上升的特征图样。那么我们要如何用代码描绘这个过程呢？作者公开的源代码给予了我们一种思路，其主要应用了中的
函数。先简单介绍一下这个函数。
是一个用于索引给定张量中某一个维度中某些特定索引元素的方法，其API手册如：
torch.index_select(input, dim, index, out=None) → Tensor
Parameters:
input (Tensor) – 输入张量，需要被索引的张量
dim (int) – 在某个维度被索引
index (LongTensor) – 一维张量，用于提供索引信息
out (Tensor, optional) – 输出张量，可以不填其作用很简单，比如我现在的输入张量为的尺寸大小，其中为样本数量，
为特征数目，如果我现在需要指定的某些样本，比如第
,
等等样本，我可以用一个
进行索引，然后应用
就可以索引了，例子如：
>>> x = torch.randn(3, 4)
>>> x
tensor([[ 0.1427, 0.0231, -0.5414, -1.0009],
[-0.4664, 0.2647, -0.1228, -1.1068],
[-1.1734, -0.6571, 0.7230, -0.6004]])
>>> indices = torch.tensor([0, 2])
>>> torch.index_select(x, 0, indices) # 按行索引
tensor([[ 0.1427, 0.0231, -0.5414, -1.0009],
[-1.1734, -0.6571, 0.7230, -0.6004]])
>>> torch.index_select(x, 1, indices) # 按列索引
tensor([[ 0.1427, -0.5414],
[-0.4664, -0.1228],
[-1.1734, 0.7230]])
注意到有一个问题是，似乎在使用
的情况下，不检查
是否会越界，因此如果你的
越界了，但是报错的地方可能不在使用
的地方，而是在后续的代码中，这个似乎就需要留意下你的
了。同时，
是一个
，这个也是要留意的。
我们先贴出主要代码，看看作者是怎么实现的：
class Shift_gcn(nn.Module):
def __init__(self, in_channels, out_channels, A, coff_embedding=4, num_subset=3):
super(Shift_gcn, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
if in_channels != out_channels:
self.down = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 1),
nn.BatchNorm2d(out_channels)
)
else:
self.down = lambda x: x
self.Linear_weight = nn.Parameter(torch.zeros(in_channels, out_channels, requires_grad=True), requires_grad=True)
self.Linear_bias = nn.Parameter(torch.zeros(1,1,out_channels,requires_grad=True),requires_grad=True)
self.Feature_Mask = nn.Parameter(torch.ones(1,25,in_channels, requires_grad=True),requires_grad=True)
self.bn = nn.BatchNorm1d(25*out_channels)
self.relu = nn.ReLU()
index_array = np.empty(25*in_channels).astype(np.int)
for i in range(25):
for j in range(in_channels):
index_array[i*in_channels + j] = (i*in_channels + j + j*in_channels) % (in_channels*25)
self.shift_in = nn.Parameter(torch.from_numpy(index_array),requires_grad=False)
index_array = np.empty(25*out_channels).astype(np.int)
for i in range(25):
for j in range(out_channels):
index_array[i*out_channels + j] = (i*out_channels + j - j*out_channels) % (out_channels*25)
self.shift_out = nn.Parameter(torch.from_numpy(index_array),requires_grad=False)
def forward(self, x0):
n, c, t, v = x0.size()
x = x0.permute(0,2,3,1).contiguous()
# n,t,v,c
# shift1
x = x.view(n*t,v*c)
x = torch.index_select(x, 1, self.shift_in)
x = x.view(n*t,v,c)
x = x * (torch.tanh(self.Feature_Mask)+1)
x = torch.einsum('nwc,cd->nwd', (x, self.Linear_weight)).contiguous() # nt,v,c
x = x + self.Linear_bias
# shift2
x = x.view(n*t,-1)
x = torch.index_select(x, 1, self.shift_out)
x = self.bn(x)
x = x.view(n,t,v,self.out_channels).permute(0,3,1,2) # n,c,t,v
x = x + self.down(x0)
x = self.relu(x)
# print(self.Feature_Mask.shape)
return x我们把forward()里面的分为三大部分，分别是：1> shift_in操作；2> 卷积操作；3> shift_out操作；其中指的shift_in和shift_out只是shift图卷积算子的不同形式而已，其主要是一致的。整个结构图如Fig 2（c）所示。
x = torch.einsum('nwc,cd->nwd', (x, self.Linear_weight)).contiguous() # nt,v,c
x = x + self.Linear_bias
x = x * (torch.tanh(self.Feature_Mask)+1)
那么我们着重考虑以下的代码：
x = x.view(n*t,v*c)
x = torch.index_select(x, 1, self.shift_in)
x = x.view(n*t,v,c)第一行代码将特征图展开，如Fig 3所示，得到了25 × C 25 \times C25×C大小的特征向量。通过torch.index_select对特征向量的不同分区进行选择得到最终的输出特征向量，选择的过程如Fig 4所示。
那么可以知道，对于某个关节点i ii而言，给定通道j jj，当遍历不同通道时，会存在一个周期，因此是
，比如对于第0号节点的第1个通道，其需要将
的值移入，如Fig 4的例子所示。而第2个通道则是需要考虑将
的值移入，我们发现是以C CC为周期的。这个时候假定的是关节点都是同一个的时候，当遍历关节点时，我们最终的索引规则是(
，因为考虑到了溢出的问题，因此需要求余，有
。这个对应源代码的第23-32行，如上所示。
在以这个举个代码例子，例子如下所示：
import numpy as np
import torch
array = np.arange(0,15).reshape(3,5)
array = torch.tensor(array)
index = np.zeros(15)
for i in range(3):
for j in range(5):
index[i*5+j] = (i*5+j*5+j) % (15)
index = torch.tensor(index).long()
out = torch.index_select(array.view(1,-1), 1, index).view(3,5)
print(array)
print(out)输出为：
tensor([[ 0, 1, 2, 3, 4],
[ 5, 6, 7, 8, 9],
[10, 11, 12, 13, 14]])
tensor([[ 0, 6, 12, 3, 9],
[ 5, 11, 2, 8, 14],
[10, 1, 7, 13, 4]])我们把这种正向移入的称之为，反过来移入则称之为
，其索引公式有一点小变化，为：
。代码例子如下：
import numpy as np
import torch
array = np.arange(0,15).reshape(3,5)
array = torch.tensor(array)
index = np.zeros(15)
for i in range(3):
for j in range(5):
index[i*5+j] = (i*5-j*5+j) % (15)
index = torch.tensor(index).long()
out = torch.index_select(array.view(1,-1), 1, index).view(3,5)
print(array)
print(out)输出为：
tensor([[ 0, 1, 2, 3, 4],
[ 5, 6, 7, 8, 9],
[10, 11, 12, 13, 14]])
tensor([[ 0, 11, 7, 3, 14],
[ 5, 1, 12, 8, 4],
[10, 6, 2, 13, 9]])输入和只是因为平移方向反过来了而已。
当然，进行了特征向量的还不够，还需要将其
回一个特征矩阵，因此会有：
x = x.view(n*t,v,c)这样的代码段出现。
Reference
[1]. https://fesian.blog.csdn.net/article/details/109563113
[2]. https://fesian.blog.csdn.net/article/details/109474701
[3]. https://github.com/kchengiva/Shift-GCN
[4]. https://blog.csdn.net/LoseInVain/article/details/81143966