模型结构：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(
	(0): Bottleneck(
	  (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	  (downsample): Sequential(
		(0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
		(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  )
	)
	(1): Bottleneck(
	  (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
	(2): Bottleneck(
	  (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
  )
  (layer2): Sequential(
	(0): Bottleneck(
	  (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	  (downsample): Sequential(
		(0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
		(1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  )
	)
	(1): Bottleneck(
	  (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
	(2): Bottleneck(
	  (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
	(3): Bottleneck(
	  (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
  )
  (layer3): Sequential(
	(0): Bottleneck(
	  (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	  (downsample): Sequential(
		(0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)
		(1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  )
	)
	(1): Bottleneck(
	  (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
	(2): Bottleneck(
	  (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
	(3): Bottleneck(
	  (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
	(4): Bottleneck(
	  (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
	(5): Bottleneck(
	  (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
  )
  (layer4): Sequential(
	(0): Bottleneck(
	  (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	  (downsample): Sequential(
		(0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)
		(1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  )
	)
	(1): Bottleneck(
	  (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
	(2): Bottleneck(
	  (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
	  (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
	  (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	  (relu): ReLU(inplace)
	)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): Linear(in_features=2048, out_features=1000, bias=True)
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def make_temporal_shift(net, n_segment, n_div=8, place='blockres', temporal_pool=False):
	if temporal_pool:
		n_segment_list = [n_segment, n_segment // 2, n_segment // 2, n_segment // 2]
	else:
		n_segment_list = [n_segment] * 4
	assert n_segment_list[-1] > 0
	#print('=> n_segment per stage: {}'.format(n_segment_list))

	import torchvision
	if isinstance(net, torchvision.models.ResNet):
		if place == 'block':
			def make_block_temporal(stage, this_segment):
				blocks = list(stage.children())
				#print('=> Processing stage with {} blocks'.format(len(blocks)))
				for i, b in enumerate(blocks):
					blocks[i] = TemporalShift(b, n_segment=this_segment, n_div=n_div)
				return nn.Sequential(*(blocks))

			net.layer1 = make_block_temporal(net.layer1, n_segment_list[0])
			net.layer2 = make_block_temporal(net.layer2, n_segment_list[1])
			net.layer3 = make_block_temporal(net.layer3, n_segment_list[2])
			net.layer4 = make_block_temporal(net.layer4, n_segment_list[3])

		elif 'blockres' in place:
			n_round = 1
			if len(list(net.layer3.children())) >= 23: # 如果是ResNet101的话，就每两轮采一次
				n_round = 2
				#print('=> Using n_round {} to insert temporal shift'.format(n_round))

			def make_block_temporal(stage, this_segment):
				blocks = list(stage.children())
				#print('=> Processing stage with {} blocks residual'.format(len(blocks)))
				for i, b in enumerate(blocks):
					if i % n_round == 0:
						blocks[i].conv1 = TemporalShift(b.conv1, n_segment=this_segment, n_div=n_div)
				return nn.Sequential(*blocks)

			net.layer1 = make_block_temporal(net.layer1, n_segment_list[0])
			net.layer2 = make_block_temporal(net.layer2, n_segment_list[1])
			net.layer3 = make_block_temporal(net.layer3, n_segment_list[2])
			net.layer4 = make_block_temporal(net.layer4, n_segment_list[3])
	else:
		raise NotImplementedError(place)

简单来说，就是对每一层的第一个卷积操作之前，进行shift操作。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
@staticmethod
def shift(x, n_segment, fold_div=3, inplace=False):
	nt, c, h, w = x.size()
	n_batch = nt // n_segment
	x = x.view(n_batch, n_segment, c, h, w)

	fold = c // fold_div # fold_div 为 8，论文中说他们测试发现1/4的shift，即左1/8和右1/8效果最好
	if inplace:
		# Due to some out of order error when performing parallel computing. 
		# May need to write a CUDA kernel.
		raise NotImplementedError  
		# out = InplaceShift.apply(x, fold)
	else:
		out = torch.zeros_like(x) # 生成一个和x维度信息一样，但是全是0的数组
		# 假设原来是 [1,2,3,4,5,6,7,8]
		
		out[:, :-1, :fold] = x[:, 1:, :fold]  # shift left
		# 变成了 [2,3,4,5,6,7,8,0]
		
		out[:, 1:, fold: 2 * fold] = x[:, :-1, fold: 2 * fold]  # shift right
		# 变成了 [0,1,2,3,4,5,6,7]
		
		out[:, :, 2 * fold:] = x[:, :, 2 * fold:]  # not shift
		# 其余的不发生变化
	
	return out.view(nt, c, h, w)

对通道进行分片，前1/8的通道左移，再1/8的通道右移，剩余的保持不变。