pytorch - Faster RCNN模型(Pytorch)预测结果为空

Question

• 我正在尝试训练 Faster RCNN 模型。训练后，我尝试预测图像的结果，但结果为空。
• 我的数据是 w: 1600, h: 800, c: 3, classes: 7, bounding boxes:(x1, y1, x2, y2)
• 下面是我的模型。

我的模型

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
      
def get_instance_segmentation_model(num_classes):
    backbone = torchvision.models.vgg16(pretrained=True).features
    backbone.out_channels = 512

    anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
                                    aspect_ratios=((0.5, 1.0, 2.0),))

    roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
                                                    output_size=7,
                                                    sampling_ratio=2)
    model = FasterRCNN(backbone,
                    num_classes=2,
                    rpn_anchor_generator=anchor_generator,
                    box_roi_pool=roi_pooler)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)


    return model

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

num_classes = 2

model = get_instance_segmentation_model(num_classes)

model.to(device)


params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
                            momentum=0.9, weight_decay=0.0005)

lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                               step_size=3,
                                               gamma=0.1)

训练

# let's train it for 10 epochs
num_epochs = 10

for epoch in range(num_epochs):
    # train for one epoch, printing every 10 iterations
    train_one_epoch(model, optimizer, train_data_loader, device, epoch, print_freq=10)
    # update the learning rate
    lr_scheduler.step()
    # evaluate on the test dataset
    evaluate(model, valid_data_loader, device=device)

预测:

prediction
[{'boxes': tensor([], device='cuda:0', size=(0, 4)),
  'labels': tensor([], device='cuda:0', dtype=torch.int64),
  'scores': tensor([], device='cuda:0')}]

Answer 1

您应该将类​​数更改为

    model = FasterRCNN(backbone,
                    num_classes=YOUR_CLASSES+1,  # +1 is for the background
                    rpn_anchor_generator=anchor_generator,
                    box_roi_pool=roi_pooler)

请记住类 0 是为后台保留的，因此您的类应该从 1 开始。

此外，请确保您的网络在训练期间已经收敛。