class ResNet(nn.Module):
def __init__(self):
super(ResNet, self).__init__()
resnet = models.resnet50(pretrained=True)
modules = list(resnet.children())[:-1]
self.resnet = nn.Sequential(*modules)
self.fc=nn.Linear(2048,10)
def forward(self, x):
x1 = self.resnet(x)
x1=self.resnet(x)
x1 = x1.view(x1.size(0), -1)
x1=self.fc(x1)
x2 = self.resnet(x)
x2= x2.view(x2.size(0), -1)
x2=self.fc(x2)
return x1,x2
-
$\begingroup$ Welcome to the forum! Can you add a little more description to your question. This makes it easier for others to understand the problem and increases your chance of getting a helpful reply. $\endgroup$– PeterJun 3, 2019 at 10:52
1 Answer
A softmax layer helps:
class ResNet(nn.Module): def init(self): super(ResNet, self).init() resnet = models.resnet50(pretrained=True) modules = list(resnet.children())[:-1] self.resnet = nn.Sequential(*modules) self.fc=nn.Linear(2048,10),
def forward(self, x):
x1 = torch.softmax(self.resnet(x), dim=-1)
#x1=self.resnet(x)
#x1 = x1.view(x1.size(0), -1)
#x1=self.fc(x1)
x2 = torch.softmax(self.resnet(x) dim=-1)
#x2= x2.view(x2.size(0), -1)
#x2=self.fc(x2)
return x1,x2
To know more about the dim attribute, corss reference here: https://stackoverflow.com/questions/52513802/pytorch-softmax-with-dim
I hope this what you were looking for.