pytorch使用GPU上训练

使用单GPU来训练

#使用GPU 2 来进行训练
device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")

model.to(device)

data = data.to(device)
label=label.to(device)

prediction= model(data)

###先转为cpu上的tensor 再从tensor转为numpy
prediction=prediction.cpu().numpy

测试代码

import torch
import torch.nn as nn
import torch.utils.data as Data
import  torchvision
import os


device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")




EPOCH = 5
BATCH_SIZE = 50
LR = 0.001
DOWNLOAD_MNIST = True

train_data = torchvision.datasets.MNIST(
    root='/data/hui/Minst/dataset/',
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=DOWNLOAD_MNIST,
)

test_data = torchvision.datasets.MNIST(root='/data/hui/Minst/dataset/', train=False)

# 批处理
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
# 测试
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(torch.FloatTensor)[:2000] / 255
test_y = test_data.test_labels[:2000]

test_x = test_x.to(device)
test_y = test_y.to(device)


# 卷积(Conv2d) -> 激励函数(ReLU) -> 池化, 向下采样 (MaxPooling) ->
# 再来一遍 -> 展平多维的卷积成的特征图 -> 接入全连接层 (Linear) -> 输出

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential( # 1x28x28
            nn.Conv2d(
                in_channels=1,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2,
            ),# 16x28x28
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.conv2 = nn.Sequential( # 16x14x14
            nn.Conv2d(16, 32, 5, 1, 2),# 32x14x14
            nn.ReLU(),
            nn.MaxPool2d(2),# 32x7x7
        )
        self.out = nn.Linear(32*7*7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)  # 展平多维的卷积图成 (batch_size, 32 * 7 * 7)\
        output = self.out(x)
        return output

cnn = CNN()

cnn = cnn.to(device)



print(cnn)

optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)
loss_func = nn.CrossEntropyLoss()

for epoch in range(EPOCH):
    for step, (b_x, b_y) in enumerate(train_loader):
        
        output = cnn(b_x.to(device))
        loss = loss_func(output, b_y.to(device))
        

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

test_output = cnn(test_x[:10])
pred_y = torch.max(test_output, 1)[1].data.cpu().numpy().squeeze()
print(pred_y, 'prediction number')
print(test_y[:10].cpu().numpy(), 'real number')

或者 另外一种方法

model = model.cuda(device=3)  


data = data.cuda(device=3)
label=label.cuda(device=3)

prediction= model(data)

  
  

使用多GPU训练

##使用GPU 2 3来进行训练
device_ids = [2,3]
model = torch.nn.DataParallel(model, device_ids=device_ids) # 声明所有可用设备
model = cnn.cuda(device=device_ids[0])  # 模型放在主设备

#将数据和标签都放在主设备上
data=data.cuda(device=device_ids[0])  
label=label.cuda(device=device_ids[0])

prediction= model(data)

###先转为cpu上的tensor 再从tensor转为numpy
prediction=prediction.cpu().numpy

测试代码

import torch
import torch.nn as nn
import torch.utils.data as Data
import  torchvision
import os

device_ids = [2,3]


EPOCH = 5
BATCH_SIZE = 50
LR = 0.001
DOWNLOAD_MNIST = True

train_data = torchvision.datasets.MNIST(
    root='/data/chenxiaohui/hui/dataset/',
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=DOWNLOAD_MNIST,
)

test_data = torchvision.datasets.MNIST(root='/data/hui/Minst/dataset/', train=False)

# 批处理
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
# 测试
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(torch.FloatTensor)[:2000] / 255
test_y = test_data.test_labels[:2000]

# test_x = test_x.to(device)
# test_y = test_y.to(device)


# 卷积(Conv2d) -> 激励函数(ReLU) -> 池化, 向下采样 (MaxPooling) ->
# 再来一遍 -> 展平多维的卷积成的特征图 -> 接入全连接层 (Linear) -> 输出

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential( # 1x28x28
            nn.Conv2d(
                in_channels=1,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2,
            ),# 16x28x28
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.conv2 = nn.Sequential( # 16x14x14
            nn.Conv2d(16, 32, 5, 1, 2),# 32x14x14
            nn.ReLU(),
            nn.MaxPool2d(2),# 32x7x7
        )
        self.out = nn.Linear(32*7*7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)  # 展平多维的卷积图成 (batch_size, 32 * 7 * 7)\
        output = self.out(x)
        return output

cnn = CNN()

cnn = torch.nn.DataParallel(cnn, device_ids=device_ids) # 声明所有可用设备
cnn = cnn.cuda(device=device_ids[0])  # 模型放在主设备

print(cnn)

optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)
loss_func = nn.CrossEntropyLoss()

for epoch in range(EPOCH):
    for step, (b_x, b_y) in enumerate(train_loader):
        

        
        
        output = cnn(b_x.cuda(device=device_ids[0]))
        loss = loss_func(output, b_y.cuda(device=device_ids[0]))
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

test_output = cnn(test_x[:10])
pred_y = torch.max(test_output, 1)[1].data.cpu().numpy().squeeze()
print(pred_y, 'prediction number')
print(test_y[:10].cpu().numpy(), 'real number')