if __name__ == '__main__':
import torch.nn as nn
import torch
x=[torch.randn(1,256,256,256),torch.randn(1,512,128,128),torch.randn(1,1024,64,64),torch.randn(1,2048,32,32)]
mode = YOLOXPAFPN(in_channels=[256, 512, 1024, 2048],
out_channels=256,
num_csp_blocks=1)
outputs = mode.forward(x)
for i in range(len(outputs)):
print(f'outputs[{i}].shape = {outputs[i].shape}')
outputs[0].shape = torch.Size([1, 256, 256, 256])
outputs[1].shape = torch.Size([1, 256, 128, 128])
outputs[2].shape = torch.Size([1, 256, 64, 64])
outputs[3].shape = torch.Size([1, 256, 32, 32])
def forward(self, inputs):
"""
Args:
inputs (tuple[Tensor]): input features.
Returns:
tuple[Tensor]: YOLOXPAFPN features.
"""
assert len(inputs) == len(self.in_channels)
# top-down path
inner_outs = [inputs[-1]]
for idx in range(len(self.in_channels) - 1, 0, -1):
feat_heigh = inner_outs[0]
feat_low = inputs[idx - 1]
feat_heigh = self.reduce_layers[len(self.in_channels) - 1 - idx](
feat_heigh)
inner_outs[0] = feat_heigh
upsample_feat = self.upsample(feat_heigh)
inner_out = self.top_down_blocks[len(self.in_channels) - 1 - idx](
torch.cat([upsample_feat, feat_low], 1))
inner_outs.insert(0, inner_out)
# bottom-up path
outs = [inner_outs[0]]
for idx in range(len(self.in_channels) - 1):
feat_low = outs[-1]
feat_height = inner_outs[idx + 1]
downsample_feat = self.downsamples[idx](feat_low)
out = self.bottom_up_blocks[idx](
torch.cat([downsample_feat, feat_height], 1))
outs.append(out)
# out convs
for idx, conv in enumerate(self.out_convs):
outs[idx] = conv(outs[idx])
return tuple(outs)
分为三部分,第一部分从上到下,第二部分从下到上,第三部分改变输出通道数
inner out 为2048 32 32
第一个循环
2048*32*32,将通道 变为 1024*32*32
inner_outs[0] = feat_heigh, inner out 为1024 32 32
然后上采样 特征图变为1024*64*64
和1024*64*64 合并变为2048*64*64
进入csp通道数变为1024*64*64
inner outs 为
1024*64*64
1024 32 32
====================
第二个循环
1024*64*64,将通道 变为 512*64*64
inner_outs[0] = 512*64*64
然后上采样 特征图变为512*128*128
和512*128*128合并变为1024*128*128
进入csp通道数变为512*128*128
inner outs 为
512*128*128
512*64*64
1024 32 32
====================
第三个循环
512*128*128,将通道 变为 256*128*128
inner_outs[0] = 256*128*128
然后上采样 特征图变为256*256*256
和256*256*256合并变为512*256*256
进入csp通道数变为256*256*256
inner outs 为
256*256*256
256*128*128
512*64*64
1024 32 32
第一部分结束
------------------------------------------------------
第二部分
inner outs中,
第一个循环
256*256*256 降采样 256*128*128
合并变为 512*128*128
经过csp变为 512*128*128
out为256*256*256
512*128*128
============================
第二个循环
512*128*128降采样 512*64*64
合并变为 1024*64*64
经过csp变为 1024*64*64
out为256*256*256
512*128*128
1024*64*64
============================
第三个循环
1024*64*64降采样 1024*32*32
合并变为 2048*32*32
经过csp变为 2048*32*32
out为 256*256*256
512*128*128
1024*64*64
2048*32*32
--------------------------------------
第三部分,把四个输出的通道数都变成256
输入(1,256,256,256),(1,512,128,128),(1,1024,64,64),(1,2048,32,32)
总结,第一部分,小特征图,降通道,将通道结果保存在inner_out 中,上采样,与下层通道数和大小都一致,然后合并通道会加倍,然后在CSP中通道减半。
此时输出为
256*256*256
256*128*128
512*64*64
1024 32 32 下面三层的通道数都少了一半
第二部分,大特征图下采样,合并通道加倍,经过csp,但是不改变特征图通道数。
输出为
256*256*256
512*128*128
1024*64*64
2048*32*32