我可以使用预训练的 YOLO 进行回归任务吗？

Question

我计划使用 YOLO 来完成 CNN 监督回归任务。给定一张图像，预测该图像将被查看的次数。我倾向于使用 YOLO，因为它是一个物体检测器。观看次数较多的照片大多包含对象（面部、动物、文本等），这些对象是 YOLO 最初训练的 COCO 数据集中的类。

我已经尝试过使用预训练的 CNN 模型（VGGNet、MobileNet 等）和冻结权重，但结果并不好。微调预训练模型的选项是不可能的，因为我没有计算资源来使用 100K+ 图像进行 x epoch 的训练，只是为了为我的问题创建一个好的模型。

Answer 1

YOLO

使用 Darknet 作为 CNN 主干/特征提取器。因此，您可能想尝试使用预先训练的暗网作为特征提取器，并用回归器替换分类器。标准 YOLO 使用 Darknet-53，而 Tiny YOLO 使用 Darknet-19。

Answer 2

你应该做的是让模型架构中的最后一层接受回归数据

示例：

self.vit_b_16.heads = nn.Linear(self.vit_b_16.heads[0].in_features, num_classes)

然后为您的模型创建一个数据加载器来处理它

from torch.utils.data import DataLoader, Dataset

class RegressionDataset(Dataset):
    def __init__(self, image_folder, targets_csv, transform=None):
        self.image_folder = image_folder
        self.targets_df = pd.read_csv(targets_csv)
        self.transform = transform

    def __len__(self):
        return len(self.targets_df)

    def __getitem__(self, idx):
        img_name = os.path.join(self.image_folder, self.targets_df.iloc[idx, 0])
        image = Image.open(img_name).convert('RGB')

        if self.transform:
            image = self.transform(image)

        target = torch.tensor(self.targets_df.iloc[idx, 1], dtype=torch.float32)

        return image, target