submodule.py
class feature_extraction(nn.Module):
def __init__(self):
print('feature_extraction 123')
super(feature_extraction, self).__init__()
self.inplanes = 32
self.firstconv = nn.Sequential(convbn(3, 32, 3, 2, 1, 1),
nn.ReLU(inplace=True),
convbn(32, 32, 3, 1, 1, 1),
nn.ReLU(inplace=True),
convbn(32, 32, 3, 1, 1, 1),
nn.ReLU(inplace=True))
self.layer1 = self._make_layer(BasicBlock, 32, 3, 1,1,1)
self.layer2 = self._make_layer(BasicBlock, 64, 16, 2,1,1)
self.layer3 = self._make_layer(BasicBlock, 128, 3, 1,1,1)
self.layer4 = self._make_layer(BasicBlock, 128, 3, 1,1,2)
def forward(self, x):
print('feature_extraction forward 123')
output = self.firstconv(x)
output = self.layer1(output)
output_raw = self.layer2(output)
output = self.layer3(output_raw)
output_skip = self.layer4(output)
对应下面
