Upload folder using huggingface_hub

model.py

@@ -0,0 +1,35 @@
+import torch.nn as nn
+import torch
+from torchvision.models import resnet50
+import os
+os.environ["TORCH_HOME"] = os.getcwd()
+
+# reference : https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py#L166
+
+class Resnet50Custom(nn.Module):
+    def __init__(self,  object_type_output_dim=15):
+        super().__init__()
+        self.resnet50 = resnet = resnet50(weights="IMAGENET1K_V2")
+        for param in self.resnet50.parameters():
+            param.requires_grad = False
+        self.resnet50.fc = nn.Linear(self.resnet50.fc.in_features, object_type_output_dim)
+        self.ff_defect_classification = nn.Linear(self.resnet50.fc.in_features, 2)
+
+    def forward(self, x):
+        x = self.resnet50.conv1(x)
+        x = self.resnet50.bn1(x)
+        x = self.resnet50.relu(x)
+        x = self.resnet50.maxpool(x)
+
+        x = self.resnet50.layer1(x)
+        x = self.resnet50.layer2(x)
+        x = self.resnet50.layer3(x)
+        x = self.resnet50.layer4(x)
+
+        x = self.resnet50.avgpool(x)
+        x = torch.flatten(x, 1)
+
+        ff_object_classification_output = self.resnet50.fc(x)
+        ff_defect_classification_output = self.ff_defect_classification(x)
+
+        return ff_object_classification_output, ff_defect_classification_output