could anyone help me solve this problem,
code below:
import torch
from torch import nn, optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
import torchvision.models as models
from torchvision import datasets, transforms
class CNN(nn.Module):
def init(self,size):
super(CNN, self).__init__()
resnet = models.resnet152(pretrained=True)
modules = list(resnet.children())[:-1] # delete the last fc layer.
self.resnet = nn.Sequential(*modules)
self.linear = nn.Linear(resnet.fc.in_features, size)
self.bn = nn.BatchNorm1d(size, momentum=0.01)
def forward(self, images):
"""Extract feature vectors from input images."""
with torch.no_grad():
features = self.resnet(images)
features = features.reshape(features.size(0), -1)
features = self.bn(self.linear(features))
return features
device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)
NET = CNN().cuda()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(NET.parameters(), lr=0.01)
for epoch in range (10):
for img, rst1, rst2 in train_dataloader:
img=img.to(device)
rst1 = rst1.to(device)
rst2 = rst2.to(device)
outputs = NET(img)
loss = criterion(outputs, (rst1,rst2))
optimizer.zero_grad()
loss.backward()
optimizer.step()