add as_sequential func

Wuhall · Dec 21, 2022 · b326220 · b326220
1 parent dfef4da
commit b326220
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/model/backbone/MobileNetV3.py b/model/backbone/MobileNetV3.py
@@ -153,3 +153,10 @@ def __init__(
         self.classifier = Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
 
         efficientnet_init_weights(self)
+
+    def as_sequential(self):
+        layers = [self.conv_stem, self.bn1]
+        layers.extend(self.blocks)
+        layers.extend([self.global_pool, self.conv_head, self.act2])
+        layers.extend([nn.Flatten(), nn.Dropout(self.drop_rate), self.classifier])
+        return nn.Sequential(*layers)