R-CNN、Fast R-CNN、Faster R-CNN ——简单实现

R-CNN

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import numpy as np

class RCNN(nn.Module):
    def __init__(self, num_classes):
        super(RCNN, self).__init__()
        # 使用预训练的 ResNet-50 作为基础网络
        self.backbone = torchvision.models.resnet50(pretrained=True)
        # 替换 ResNet-50 的最后一层全连接层
        self.backbone.fc = nn.Linear(2048, 512)
        # 分类器和回归器
        self.classifier = nn.Linear(512, num_classes)
        self.regessor = nn.Linear(512, 4)  # 4 维向量表示边界框坐标

    def forward(self, x):
        # 基础网络前向传播
        features = self.backbone(x)
        # flatten
        features = features.view(features.size(0), -1)
        # 分类和回归
        cls_output = self.classifier(features)
        reg_output = self.regessor(features)
        return cls_output, reg_output

# 定义损失函数
class RCNNLoss(nn.Module):
    def __init__(self):
        super(RCNNLoss, self).__init__()
        self.cls_loss = nn.CrossEntropyLoss()  # 分类损失
        self.reg_loss = nn.SmoothL1Loss()      # 回归损失

    def forward(self, cls_output, reg_output, cls_target, reg_target):
        # 分类损失
        cls_loss = self.cls_loss(cls_output, cls_target)
        # 回归损失
        reg_loss = self.reg_loss(reg_output, reg_target)
        # 总损失
        total_loss = cls_loss + reg_loss
        return total_loss

# 使用示例
if __name__ == "__main__":
    # 定义模型和损失函数
    num_classes = 10
    model = RCNN(num_classes)
    criterion = RCNNLoss()

    # 假设输入图像 img 是一个 numpy 数组
    img = np.random.randn(3, 224, 224)  # 示例输入图像
    img_tensor = torch.tensor(img).unsqueeze(0)  # 增加 batch 维度

    # 假设分类和回归的目标分别是 cls_target 和 reg_target
    cls_target = torch.randint(0, num_classes, (1,)