训练更快的 RCNN 时将所有张量保留在 GPU 中

Question

我正在尝试训练一个更快的 RCNN 模型，用于在自定义的类 COCO 数据集上进行边界框检测。我正在使用 GPU，即使我使用 .to(device) 将张量推入 GPU，我仍然收到以下错误：

/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in linear(input, weight, bias)
   1845     if has_torch_function_variadic(input, weight):
   1846         return handle_torch_function(linear, (input, weight), input, weight, bias=bias)
-> 1847     return torch._C._nn.linear(input, weight, bias)
   1848 
   1849 

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking arugment for argument mat1 in method wrapper_addmm)

训练 sn-p：

# Initialize Dataset
train_dataset = TrainDataset('coco_train.json')

def collate_fn(batch):
    return tuple(zip(*batch))

train_data_loader = DataLoader(
    train_dataset,
    batch_size=2,
    shuffle=True,
    num_workers=2,
    collate_fn=collate_fn
)

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

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
model.to(device) # EDIT


num_classes = 3 # eyelids, iris + background

in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

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)
num_epochs = 40

itr = 1

for epoch in range(num_epochs):
    for images, targets in train_data_loader:

        images = list(image.to(device) for image in images)
        targets = [ { k: v.to(device) for k, v in t.items() } for t in targets ]

        loss_dict = model(images, targets)

        losses = sum(loss for loss in loss_dict.values())
        loss_value = losses.item()

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

        if itr % 50 == 0:
            print(f"Iteration #{itr} loss: {loss_value}")
        
        itr += 1
        lr_scheduler.step()
    

我的数据集getitem() sn-p:（截断）

def __getitem__(self, index : int):
    ...
    
    target = {}

    target["boxes"] = torch.as_tensor(boxes, dtype=torch.float32)
    target["labels"] = torch.as_tensor(labels, dtype=torch.int64)
    target["image_id"] = torch.as_tensor([ int(image_id) ], dtype=torch.int64)
    target["area"] = torch.as_tensor(area, dtype=torch.float32)
    target["iscrowd"] = torch.as_tensor(iscrowd, dtype=torch.int64)

    image = torchvision.transforms.ToTensor()(image)
    return image, target

Answer 1

您也应该在 GPU 设备上加载您的模型：

model.to(device)

注意，torch.nn.Module.to 是就地操作。