| import torch.nn as nn |
| import torch |
| from torchvision.models import resnet50 |
| import os |
| os.environ["TORCH_HOME"] = os.getcwd() |
|
|
| |
|
|
| 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 |
|
|
