Настройка более быстрой модели R-CNN на PyTorch с помощью магистрали Swin Transformer V2 ⇐ Python

[Anonymous]

Для моего проекта по обнаружению объектов я некоторое время использовал модель fastrcnn_resnet50_fpn_v2. Мои входные изображения имеют высокое разрешение (около 3000 x 4000 пикселей), и я объединяю их в плитки размером 1200 x 1600 пикселей для обучения и вывода. Однако я изо всех сил пытаюсь эффективно обнаруживать небольшие объекты (размером всего 10 x 10 пикселей) с помощью этой модели.
Когда искал альтернативы, я прочитал о SwinTransformer V2 и Я нашел это многообещающим, особенно для приложений с изображениями высокого разрешения. Поскольку все изображения в моем наборе данных также имеют размер 1200 x 1600, и я не хочу их уменьшать, я хотел бы настроить Faster R-CNN для использования магистральной сети Swin V2 и, возможно, добавить FPN и реализовать Cascade R-CNN. голова. Однако я сталкиваюсь с проблемами несоответствия размеров позвоночника, шеи и головы RPN.
Это то, что я придумал на данный момент (я решил использовать базовую модель );

Код: Выделить всё

import torch
from torch import nn
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.models.swin_transformer import swin_v2_b, Swin_V2_B_Weights
from torchvision.ops import MultiScaleRoIAlign
import torchvision.transforms as transforms
import requests
from PIL import Image

NUM_CLASSES = 100
trainable_layers = 2

class CustomSwin(nn.Module):
def __init__(self, backbone):
super().__init__()
self.backbone = backbone
self.out_channels = 1024

def forward(self, x):
return torch.permute(self.backbone(x), (0, 3, 1, 2))

backbone = swin_v2_b(weights=Swin_V2_B_Weights.DEFAULT)

# Remove the classification head
backbone.norm = nn.Identity()
backbone.permute = nn.Identity()
backbone.avgpool = nn.Identity()
backbone.flatten = nn.Identity()
backbone.head = nn.Identity()

# Freeze all parameters
for param in backbone.parameters():
param.requires_grad = False

# Unfreeze the last trainable_layers
for layer in list(backbone.features)[-trainable_layers:]:
for param in layer.parameters():
param.requires_grad = True

custom_backbone = CustomSwin(backbone)

# Added smaller sizes for really small objects
anchor_generator = AnchorGenerator(
sizes=((8, 16, 32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),)
)

roi_pooler = MultiScaleRoIAlign(featmap_names=["0"], output_size=7, sampling_ratio=2)

model = FasterRCNN(
custom_backbone,
num_classes=NUM_CLASSES,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler,
min_size=1224,
max_size=1632,
)

Я не уверен, является ли текущая реализация оптимальной или добавление таких компонентов, как FPN (Feature Pyramid Network) или Cascade R-CNN, повысит производительность модели (у меня довольно большой набор данных) . Кто-нибудь успешно реализовал эти модификации? Будем очень признательны за любые рекомендации!

Подробнее здесь: https://stackoverflow.com/questions/790 ... 2-backbone