| 8 | from torch import nn |
| 9 | |
| 10 | class Backbone(nn.Sequential): |
| 11 | def __init__(self, resnet): |
| 12 | super(Backbone, self).__init__( |
| 13 | OrderedDict( |
| 14 | [ |
| 15 | ["conv1", resnet.conv1], |
| 16 | ["bn1", resnet.bn1], |
| 17 | ["relu", resnet.relu], |
| 18 | ["maxpool", resnet.maxpool], |
| 19 | ["layer1", resnet.layer1], # res2 |
| 20 | ["layer2", resnet.layer2], # res3 |
| 21 | ["layer3", resnet.layer3], # res4 |
| 22 | ] |
| 23 | ) |
| 24 | ) |
| 25 | self.out_channels = 1024 |
| 26 | |
| 27 | def forward(self, x): |
| 28 | # using the forward method from nn.Sequential |
| 29 | feat = super(Backbone, self).forward(x) |
| 30 | return OrderedDict([["feat_res4", feat]]) |
| 31 | |
| 32 | |
| 33 | class Res5Head(nn.Sequential): |