Update model/backbones/mobilenetv2.py

model/backbones/mobilenetv2.py

@@ -4,18 +4,28 @@ from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
 class MobileNetV2Backbone(nn.Module):
     def __init__(self):
         super(MobileNetV2Backbone, self).__init__()
-        # Используем актуальный способ загрузки предобученной модели
-        self.model = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).features
-        self.enc_channels = [24, 32, 96, 320]  # Каналы с разных уровней feature extractor
+        backbone = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
+        features = backbone.features
+
+        # Слои, откуда берём фичи (по оригинальной архитектуре MODNet)
+        self.enc2x = features[:4]    # Примерно до блока 3
+        self.enc4x = features[4:7]   # Примерно до блока 6
+        self.enc8x = features[7:14]  # Примерно до блока 13
+        self.enc32x = features[14:] # Оставшиеся блоки
+
+        self.enc_channels = [24, 32, 96, 320]
 
     def forward(self, x):
-        return self.model(x)
+        out2x = self.enc2x(x)
+        out4x = self.enc4x(out2x)
+        out8x = self.enc8x(out4x)
+        out32x = self.enc32x(out8x)
+        return [out2x, out4x, out8x, out32x]
 
 
-# Добавим SUPPORTED_BACKBONES прямо сюда (или перенеси в backbones/__init__.py, если нужно)
 SUPPORTED_BACKBONES = {
     'mobilenetv2': lambda _=None: MobileNetV2Backbone()
-}  # _=None позволяет игнорировать входной параметр
+}
 
 
 def load_modnet_checkpoint(modnet, path, device):
@@ -23,10 +33,10 @@ def load_modnet_checkpoint(modnet, path, device):
     new_state_dict = {}
     for k, v in state_dict.items():
         if k.startswith('module.'):
-            k = k[len('module.'):]
+            k = k[len('module.'):]  # удаляем префикс, если обучалось с DataParallel
         new_state_dict[k] = v
 
     missing, unexpected = modnet.load_state_dict(new_state_dict, strict=False)
     print("MISSING:", missing)
     print("UNEXPECTED:", unexpected)
-    return modnet
+    return modnet