def __init__(self, num_class=2):
super(FusionNet, self).__init__()
self.color_moudle = Net(num_class=num_class, is_first_bn=True)
self.depth_moudle = Net(num_class=num_class, is_first_bn=True)
self.ir_moudle = Net(num_class=num_class, is_first_bn=True)
self.color_SE = SEModule(512, reduction=16)
self.depth_SE = SEModule(512, reduction=16)
self.ir_SE = SEModule(512, reduction=16)
self.bottleneck = nn.Sequential(
nn.Conv2d(512 * 3, 512, kernel_size=1, padding=0),
nn.BatchNorm2d(512), nn.ReLU(inplace=True))
self.res_0 = self._make_layer(SEResNeXtBottleneck,
planes=256,
blocks=2,
stride=2,
groups=32,
reduction=16,
downsample_kernel_size=1,
downsample_padding=0)
self.res_1 = self._make_layer(SEResNeXtBottleneck,
planes=512,
blocks=2,
stride=2,
groups=32,
reduction=16,
downsample_kernel_size=1,
downsample_padding=0)
self.fc = nn.Sequential(nn.Dropout(0.5), nn.Linear(2048, 256),
nn.ReLU(inplace=True),
nn.Linear(256, num_class))
def run_check_net():
num_class = 2
net = Net(num_class)