File size: 1,260 Bytes
8146246 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | 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
|