machine-learning - RuntimeError：输入类型（torch.cuda.FloatTensor）和权重类型（torch.FloatTensor）应该相同

Question

我将模型和数据设置为同一设备，但总是引发如下错误：

RuntimeError：输入类型（torch.cuda.FloatTensor）和权重类型（torch.FloatTensor）应该相同

以下是我的训练代码，希望你能解答。谢谢！

def train(train_img_path, train_label_path, pths_path, interval, log_file):

    file_num = len(os.listdir(train_img_path))
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    net = EAST(extractor=extractor, geometry_mode=geometry_mode, pretrained=True)
    net = net.to(device)
    trainset = custom_dataset(train_img_path, train_label_path)
    train_loader = data.DataLoader(trainset, batch_size=batch_size,
                                   shuffle=True, num_workers=num_workers, drop_last=True)


    optimizer = optim.SGD(net.parameters(), lr=initial_lr, momentum=momentum, weight_decay=weight_decay_sgd)
    criterion = Loss(weight_geo, weight_angle, geometry_mode="RBOX")

    net.train()
    epoch_loss = 0.
    for epoch in range(max_epoch):
        epoch_time = time.time()
        for i, (img, score_gt, geo_gt, ignored_map) in enumerate(train_loader):
            start_time = time.time()
            img, score_gt, geo_gt, ignored_map = img.to(device), score_gt.to(device),\
                                                 geo_gt.to(device), ignored_map.to(device)

            score_pred, geo_pred = net(img)
            total_loss, score_loss, loss_AABB, loss_angle = criterion(score_pred, geo_pred, score_gt, geo_gt, ignored_map)


            epoch_loss += total_loss.item()

            optimizer.zero_grad()
            total_loss.backward()

            optimizer.step()

标签： machine-learningcomputer-visionpytorchgpu

Answer 1

我怀疑你的损失函数有它自己的一些内部参数，因此你也应该

criterion = Loss(weight_geo, weight_angle, geometry_mode="RBOX").to(device)

如果您提供完整的跟踪信息，指出究竟是哪一行导致了错误，则更容易发现错误。