Upload smp_deeplabv3plus.py

NETS/smp_deeplabv3plus.py

@@ -0,0 +1,17 @@
+from torch import nn
+import segmentation_models_pytorch as smp
+
+
+class Net(nn.Module):
+    def __init__(self, class_num, in_channels=4, encoder_name="resnet34"):
+        super(Net, self).__init__()
+        self.net = smp.deeplabv3.DeepLabV3Plus(
+            in_channels=in_channels,
+            classes=class_num,
+            encoder_name=encoder_name,
+        )
+
+    def forward(self, x):
+        x = self.net(x)
+
+        return x