在PyTorch的GPU上使用DataLoader训练CNN模型

时间:2020-08-10 11:17:56

标签: python machine-learning deep-learning pytorch

我正在尝试将胸部X射线分为两类:“正常”和“肺炎”。我的训练和测试集是带有num_workers=0, pin_memory=True的DataLoader对象。 Cuda在我的设备上可用(GTX 1060 6GB)。创建CNN之后,我致电 model = CNN().cuda()。当我尝试训练模型时,我得到RuntimeError: Expected object of backend CPU but got backend CUDA for argument #2 'weight'。为了在GPU上训练模型,我必须更改什么?

代码:

root = 'chest_xray/'

train_data = datasets.ImageFolder(os.path.join(root, 'train'), transform=train_transform)
test_data = datasets.ImageFolder(os.path.join(root, 'test'), transform=test_transform)

train_loader = DataLoader(train_data,batch_size=10,shuffle=True,num_workers=0,pin_memory=True)
test_loader = DataLoader(test_data,batch_size=10,shuffle=False,num_workers=0,pin_memory=True)

class_names = train_data.classes

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 12, 5, 1)
        self.conv2 = nn.Conv2d(12, 24, 5, 1)
        self.conv3 = nn.Conv2d(24, 30, 5, 1) 
        self.conv4 = nn.Conv2d(30, 36, 5, 1)
        self.fc1 = nn.Linear(58*58*36, 256)
        self.fc2 = nn.Linear(256, 144)
        self.fc3 = nn.Linear(144, 84)
        self.fc4 = nn.Linear(84, 16)
        self.fc5 = nn.Linear(16, 2)
        
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv3(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv4(x))
        x = F.max_pool2d(x, 2, 2)
        
        x = x.view(-1, 58*58*36)
        
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = F.relu(self.fc4(x))
        x = self.fc5(x)
        
        return F.log_softmax(x, dim=1)

model = CNN().cuda()

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

epochs = 2

train_losses = []
train_correct = []
for i in range(epochs):
    trn_correct = 0
    tst_correct = 0
    
    for b, (X_train, y_train) in enumerate(train_loader):
        y_pred = model(X_train)
        loss = criterion(y_pred, y_train)
        
        predicted = torch.max(y_pred.data, 1)[1]
        batch_correct = (predicted == y_train).sum()
        trn_correct += batch_correct
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if b % 200 == 0:
            print(f'epoch: {i+1} batch: {b} progress: {10*b/len(train_data)} loss: {loss.item()} accuracy: {10*trn_correct/b}%')
            
    train_losses.append(loss)
    train_correct.append(trn_correct)

1 个答案:

答案 0 :(得分:1)

此行之后:

for b, (X_train, y_train) in enumerate(train_loader):

添加以下内容:

X_train, y_train = X_train.cuda(), y_train.cuda()