|
import os |
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
|
|
class Encoder(nn.Module): |
|
def __init__(self): |
|
super(Encoder, self).__init__() |
|
|
|
basemodel_name = 'tf_efficientnet_b5_ap' |
|
print('Loading base model ()...'.format(basemodel_name), end='') |
|
repo_path = os.path.join(os.path.dirname(__file__), 'efficientnet_repo') |
|
basemodel = torch.hub.load(repo_path, basemodel_name, pretrained=False, source='local') |
|
print('Done.') |
|
|
|
|
|
print('Removing last two layers (global_pool & classifier).') |
|
basemodel.global_pool = nn.Identity() |
|
basemodel.classifier = nn.Identity() |
|
|
|
self.original_model = basemodel |
|
|
|
def forward(self, x): |
|
features = [x] |
|
for k, v in self.original_model._modules.items(): |
|
if (k == 'blocks'): |
|
for ki, vi in v._modules.items(): |
|
features.append(vi(features[-1])) |
|
else: |
|
features.append(v(features[-1])) |
|
return features |
|
|
|
|
|
|
