(self)
| 5 | |
| 6 | class Encoder(nn.Module): |
| 7 | def __init__(self): |
| 8 | super(Encoder, self).__init__() |
| 9 | |
| 10 | basemodel_name = 'tf_efficientnet_b5_ap' |
| 11 | print('Loading base model ()...'.format(basemodel_name), end='') |
| 12 | basemodel = torch.hub.load('rwightman/gen-efficientnet-pytorch', basemodel_name, pretrained=True) |
| 13 | print('Done.') |
| 14 | |
| 15 | # Remove last layer |
| 16 | print('Removing last two layers (global_pool & classifier).') |
| 17 | basemodel.global_pool = nn.Identity() |
| 18 | basemodel.classifier = nn.Identity() |
| 19 | |
| 20 | self.original_model = basemodel |
| 21 | |
| 22 | def forward(self, x): |
| 23 | features = [x] |
nothing calls this directly
no outgoing calls
no test coverage detected