```python
class FasterRCNN(nn.Module):
def __init__(self):
super(FasterRCNN, self).__init__()
self.backbone = torchvision.models.vgg16(pretrained=True)
self.rpn = nn.Sequential(
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(512, 9 * 4, kernel_size=1, stride=1, padding=0),
)
self.roi_pooling = nn.AdaptiveAvgPool2d((7, 7))
self.classifier = nn.Sequential(
nn.Linear(7 * 7 * 512, 4096),
nn.ReLU(),
nn.Linear(4096, 4096),
nn.ReLU(),
nn.Linear(4096, 21),
)
def forward(self, x):
x = self.backbone.features(x)
rpn_output = self.rpn(x)
roi_boxes = self.convert_to_roi_boxes(rpn_output)
roi_features = self.roi_pooling(x, roi_boxes)
roi_features = roi_features.view(roi_features.size(0), -1)
output = self.classifier(roi_features)
return output
def convert_to_roi_boxes(self, rpn_output):
rpn_output = rpn_output.permute(0, 2, 3, 1).contiguous().view(-1, 4)
roi_boxes = torch.zeros((rpn_output.size(0), 4))
roi_boxes[:, 0] = rpn_output[:, 0] - 0.5 * rpn_output[:, 2]
roi_boxes[:, 1] = rpn_output[:, 1] - 0.5 * rpn_output[:, 3]
roi_boxes[:, 2] = rpn_output[:, 0] - 0.5 * rpn_output[:, 2]
roi_boxes[:, 3] = rpn_output[:, 1] - 0.5 * rpn_output[:, 3]
return roi_boxes
# 准备数据
transform = torchvision.transforms.Compose([
# 对每张图片进行170*170的中心裁剪
torchvision.transforms.Resize((256, 256)),
# 对每张图片进行随机翻转
torchvision.transforms.RandomHorizontalFlip(),
# 把图片变成tensor张量
torchvision.transforms.ToTensor(),
# 标准化图像的每个通道
torchvision.transforms.Normalize([0.4914, 0.4822, 0.4465],
[0.2023, 0.1994, 0.2010])
])
trainset = torchvision.datasets.ImageFolder(os.path.join(rectangle_dir, 'train'),
transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True,
num_workers=2)
# 定义损失函数和优化器
net = FasterRCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 训练模型
for epoch in range(2):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 2000 == 1999: # 每2000个小批量数据打印一次损失值
print(
'[%d,%5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
print('finished Training')
Traceback (most recent call last):
File "/home/u1/ywj/deepfashion/clothes_change_ywj/use_torch_on_deepfashion.py", line 185, in <module>
outputs = net(inputs)
File "/home/u1/miniconda3/envs/clothes_change_ywj/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/home/u1/ywj/deepfashion/clothes_change_ywj/use_torch_on_deepfashion.py", line 141, in forward
roi_features = self.roi_pooling(x, roi_boxes)
File "/home/u1/miniconda3/envs/clothes_change_ywj/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
TypeError: forward() takes 2 positional arguments but 3 were given
进程已结束,退出代码1
求助,该问题如何解决