PyTorch圖像分類

目錄

一、torch和torchvision

1、torchvision.datasets

2、torchvision.models

3、torchvision.transforms

4、torchvision.utils

二、MNIST手寫數字識別

1、獲取MNIST訓練集和測試集

2、數據裝載

3、數據預覽

4、構建卷積神經網絡模型

5、對模型進行訓練和參數優化

6、對訓練模型進行保存和加載

7、MNIST手寫數字識別完整代碼

三、CIFAR10圖像分類

1、CIFAR10數據集介紹

2、CIFAR10圖像分類實現

3、在GPU上跑神經網絡


一、torch和torchvision

PyTorch 中有兩個核心的包,分別是 torch torchvision

torch.nn 包提供了很多與實現神經網絡中的具體功能相關的類,torch.optim包中提供了非常多的可實現參數自動優化的類,torch.autograd實現自動梯度的功能等。

torchvision 包含了目前流行的數據集,模型結構和常用的圖片轉換工具,它的主要功能是實現數據的處理、導入和預覽等,所以如果需要對計算機視覺的相關問題進行處理,就可以借用在torchvision包中提供的大量的類來完成相應的工作。

1、torchvision.datasets

torchvision.datasets 中包含以下數據集:MNIST,COCO,LSUN Classification,ImageFolder,Imagenet-12,CIFAR10 and CIFAR100,STL10等。

2、torchvision.models

torchvision.models 模塊的子模塊中包含以下模型結構:AlexNet,VGG,ResNet,SqueezeNet,DenseNet

3、torchvision.transforms

(1)torchvision.transforms.Compose(transforms)

torchvision.transforms.Compose 類看作是一種容器,它能夠同時對多種數據變換進行組合。傳入的參數是一個列表,列表中的元素就是對載入的數據進行的各種變換操作。

transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize(mean=[0.5,0.5,0.5],
                                std=[0.5,0.5,0.5])])

torchvision.transforms.Compose 類中只是用了一個類型的轉換變化transfroms.ToTensor和一個數據標準化變換transforms.Normalize。這裏使用的是標準化變換也叫標準差變換法,這種方法需要使用原始數據的均值(Mean)和標準差(Standard Deviation)來進行數據的標準化,在經過標準化變換之後,數據全部符合均值爲0,標準差爲1的標準正態分佈。

(2)torchvision.transforms.Resize

用於對載入的圖片數據按我們需求的大小進行縮放。傳遞給這個類的參數可以是一個整型數據,也可以是一個類似於(h ,w )的序列,其中,h 代表高度,w 代表寬度,但是如果使用的是一個整型數據,那麼表示縮放的寬度和高度都是這個整型數據的值。

(3)torchvision.transforms.Scale

用於對載入的圖片數據按我們需求的大小進行縮放,用法和 torchvision.transforms.Resize 類似。

(4)torchvision.transforms.CenterCrop

用於對載入的圖片以圖片中心爲參考點,按我們需要的大小進行裁剪。傳遞給這個類的參數可以是一個整型數據,也可以是一個類似於(h ,w )的元組序列。

(5)torchvision.transforms.RandomCrop

用於對載入的圖片按我們需要的大小進行隨機裁剪。傳遞給這個類的參數可以是一個整型數據,也可以是一個類似於(h ,w )的元組序列。

(6)torchvision.transforms.RandomHorizontalFlip

用於對載入的圖片按隨機概率進行水平翻轉。我們可以通過傳遞給這個類的參數自定義隨機概率,如果沒有定義,則使用默認的概率值0.5。

(7)torchvision.transforms.RandomVerticalFlip

用於對載入的圖片按隨機概率進行垂直翻轉。我們可以通過傳遞給這個類的參數自定義隨機概率,如果沒有定義,則使用默認的概率值0.5。

(8)torchvision.transforms.ToTensor

用於對載入的圖片數據進行類型轉換,將之前構成PIL圖片的數據轉換成Tensor數據類型的變量,即將一個取值範圍是[0, 255]PIL.Imageshape(H, W, C)numpy.ndarray,轉換成形狀爲[C, H, W],取值範圍是[0, 1.0]torch.FloatTensor,讓PyTorch能夠對其進行計算和處理。

(9) torchvision.transforms.ToPILImage

用於將Tensor變量的數據轉換成PIL圖片數據,主要是爲了方便圖片內容的顯示。

(10)torchvision.transforms.RandomSizedCrop(size, interpolation=2)

先隨機切,再resize成給定size大小。

(11)torchvision.transforms.Pad(padding, fill=0)

給所有邊用給定的值填充。padding:要填充多少像素。

(12)torchvision.transforms.Normalize(mean, std)

給定均值與方差,正則化,即Normalized_image=(image-mean)/std

13)通用變換:使用lambda作爲轉換器,transforms.Lambda(lambda)

4、torchvision.utils

(1)torchvision.utils.make_grid

utils.make_grid(tensor, nrow=8, padding=2, normalize=False, range=None, scale_each=False), 給定4D-mini-batch Tensor,形狀爲(B*C*H*W),或者一個a list of image,做成一個size(B / nrow, nrow)的子圖集,其中參數:normalize=True, 對圖像像素歸一化,range=(min, max),min和max是數字,則min, max用來規範化image,scale_each=True, 每個圖片獨立規範化。

(2)torchvision.utils.save_image

utils.save_image(tensor, filename, nrow=8, padding=2, normalize=False, range=None, scale_each=False),將給定的Tensor保存成image文件,如果是mini-batch tensor,就用make-grid做成子圖集再保存。

二、MNIST手寫數字識別

1、獲取MNIST訓練集和測試集

# 對數據進行載入及有相應變換,將Compose看成一種容器,他能對多種數據變換進行組合
# 傳入的參數是一個列表,列表中的元素就是對載入的數據進行的各種變換操作(只有一個顏色通道)
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5,],std=[0.5,])])
# 獲取MNIST訓練集和測試集
data_train=datasets.MNIST(root='data/',transform=transform,train=True,download=True)
data_test=datasets.MNIST(root='data/',transform=transform,train=False)

其中,root 用於指定數據集在下載之後的存放路徑;transform用於指定導入數據集時需要對數據進行哪種變換操作,要提前定義這些變換操作;train用於指定在數據集下載完成後需要載入哪部分數據,如果設置爲True,則說明載入的是該數據集的訓練集部分;如果設置爲False,則說明載入的是該數據集的測試集部分。

2、數據裝載

在數據下載完成並且載入後,我們還需要對數據進行裝載。我們可以將數據的載入理解爲對圖片的處理,在處理完成後,我們就需要將這些圖片打包好送給我們的模型進行訓練了,而裝載就是這個打包的過程。在裝載時通過batch_size的值來確認每個包的大小,通過shuffle的值來確認是否在裝載的過程中打亂圖片的順序。

對數據的裝載使用的是torch.utils.data.DataLoader類,類中的dataset參數用於指定我們載入的數據集名稱,batch_size參數設置了每個包中的圖片數據個數,代碼中的值是64,所以在每個包中會包含64張圖片。將shuffle參數設置爲True,在裝載的過程會將數據隨機打亂順序並進行打包。

# 數據裝載
data_loader_train=torch.utils.data.DataLoader(dataset=data_train,batch_size=64,shuffle=True)
data_loader_test = torch.utils.data.DataLoader(dataset =data_test,batch_size = 64,shuffle = True)

3、數據預覽

#數據預覽和圖片顯示
images,labels=next(iter(data_loader_train))
img=torchvision.utils.make_grid(images)
img=img.numpy().transpose(1,2,0)
std=[0.5,0.5,0.5]
mean=[0.5,0.5,0.5]
img=img*std+mean
print([labels[i] for i in range(16)])
plt.imshow(img)
plt.show()

使用 iter next 來獲取一個批次的圖片數據和其對應的圖片標籤;

使用 torchvision.utils.make_grid 類方法將一個批次的圖片構造成網格模式。需要傳遞給它的參數就是一個批次的裝載數據,每個批次的裝載數據都是4維的,維度的構成從前往後分別爲batch_size、channel、height和weight,分別對應一個批次中的數據個數、每張圖片的色彩通道數、每張圖片的高度和寬度。在通過torchvision.utils.make_grid之後,圖片的維度變成了(channel,height,weight),這個批次的圖片全部被整合到了一起,所以在這個維度中對應的值也和之前不一樣了,但是色彩通道數保持不變。

使用 Matplotlib 將數據顯示成正常的圖片形式,則使用的數據首先必須是數組,其次這個數組的維度必須是 (height,weight,channel),即色彩通道數在最後面。所以我們要通過numpy和transpose完成原始數據類型的轉換和數據維度的交換,這樣才能夠使用Matplotlib繪製出正確的圖像。

打印輸出了這個批次中的數據的全部標籤,如下:

[tensor(5), tensor(2), tensor(1), tensor(7), tensor(8), tensor(4), tensor(2), tensor(3), tensor(3), tensor(9), tensor(2), tensor(1), tensor(6), tensor(3), tensor(2), tensor(7), tensor(8), tensor(7), tensor(4), tensor(6), tensor(7), tensor(3), tensor(6), tensor(7), tensor(4), tensor(6), tensor(4), tensor(3), tensor(8), tensor(7), tensor(2), tensor(4), tensor(3), tensor(7), tensor(0), tensor(2), tensor(1), tensor(4), tensor(1), tensor(0), tensor(5), tensor(0), tensor(6), tensor(3), tensor(5), tensor(9), tensor(8), tensor(0), tensor(9), tensor(0), tensor(8), tensor(3), tensor(8), tensor(2), tensor(0), tensor(5), tensor(7), tensor(6), tensor(9), tensor(1), tensor(6), tensor(0), tensor(2), tensor(9)]

對這個批次中的所有圖片數據進行顯示,如下:

4、構建卷積神經網絡模型

CNN一般結構如下:

  • 輸入層:用於數據輸入
  • 卷積層:使用卷積核進行特徵提取和特徵映射
  • 激勵層:由於卷積也是一種線性運算,因此需要增加非線性映射
  • 池化層:進行下采樣,對特徵圖稀疏處理,減少特徵信息的損失
  • 輸出層:用於輸出結果
  • CNN模型具體結構:卷積層,正則化層,激勵層,最大池化層,全連接層
# 構建卷積神經網絡模型
class CNN_Model(torch.nn.Module):
    def __init__(self):
        super(CNN_Model, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(stride=2,kernel_size=2))
        self.conv2=torch.nn.Sequential(
            torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            torch.nn.BatchNorm2d(128),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(stride=2, kernel_size=2))
        self.dense = torch.nn.Sequential(
            torch.nn.Linear(7 * 7 * 128, 1024),
            torch.nn.ReLU(),
            torch.nn.Dropout(p=0.5),
            torch.nn.Linear(1024, 10))
    # 前向傳播
    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        x = x2.view(-1, 7 * 7 * 128)
        x = self.dense(x)
        return x

torch.nn.Conv2d:用於搭建卷積神經網絡的卷積層,主要的輸入參數有輸入通道數、輸出通道數、卷積核大小、卷積核移動步長和Paddingde值。其中,輸入通道數的數據類型是整型,用於確定輸入數據的層數;輸出通道數的數據類型也是整型,用於確定輸出數據的層數;卷積核大小的數據類型是整型,用於確定卷積核的大小;卷積核移動步長的數據類型是整型,用於確定卷積核每次滑動的步長;Paddingde 的數據類型是整型,值爲0時表示不進行邊界像素的填充,如果值大於0,那麼增加數字所對應的邊界像素層數。

torch.nn.MaxPool2d:用於實現卷積神經網絡中的最大池化層,主要的輸入參數是池化窗口大小、池化窗口移動步長和Padding的值。同樣,池化窗口大小的數據類型是整型,用於確定池化窗口的大小。池化窗口步長的數據類型也是整型,用於確定池化窗口每次移動的步長。Padding的值和在torch.nn.Conv2d中定義的Paddingde值的用法和意義是一樣的。

torch.nn.Dropout:用於防止卷積神經網絡在訓練的過程中發生過擬合,其工作原理簡單來說就是在模型訓練的過程中,以一定的隨機概率將卷積神經網絡模型的部分參數歸零,以達到減少相鄰兩層神經連接的目的。可以對隨機概率值的大小進行設置,如果不足任何設置,我們就使用默認的概率值0.5。

x=x2.view(-1,7 * 7 * 128):對參數實現扁平化,因爲之後緊接着的就是全連接層,所以如果不進行扁平化,則全連接層的實際輸出的參數維度和其定義輸入的維度將不匹配,程序將會報錯。

5、對模型進行訓練和參數優化

# 對模型進行訓練和參數優化
cnn_model = CNN_Model()
loss_func = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(cnn_model.parameters(),lr=learning_rate)
n_epochs = 5
for epoch in range(n_epochs):
    running_loss = 0.0
    running_correct = 0.0
    print("Epoch  {}/{}".format(epoch, n_epochs))
    for data in data_loader_train:
        X_train, y_train = data
        X_train, y_train = Variable(X_train), Variable(y_train)
        outputs = cnn_model(X_train)
        _, pred = torch.max(outputs.data, 1)
        optimizer.zero_grad()
        loss = loss_func(outputs, y_train)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        running_correct += torch.sum(pred == y_train.data)
    testing_correct = 0.0
    for data in data_loader_test:
        X_test, y_test = data
        X_test, y_test = Variable(X_test), Variable(y_test)
        outputs = cnn_model(X_test)
        _, pred = torch.max(outputs, 1) #返回每一行中最大值的那個元素,且返回其索引
        testing_correct += torch.sum(pred == y_test.data)
        # print(testing_correct)
    print("Loss is :{:.4f},Train Accuracy is:{:.4f}%,Test Accuracy is:{:.4f}%".format(
        running_loss / len(data_train), 100 * running_correct / len(data_train),
        100 * testing_correct / len(data_test)))

6、對訓練模型進行保存和加載

# 保存模型
torch.save(cnn_model, 'data/cnn_model.pt')
# 加載模型
cnn_model=torch.load('data/cnn_model.pt')
cnn_model.eval()

7、MNIST手寫數字識別完整代碼

import torch
import torchvision
import matplotlib.pyplot as plt
from torchvision import datasets
from torchvision import transforms
from torch.autograd import Variable
# 參數設置
num_epochs = 10
batch_size = 64
learning_rate = 0.001

# 將數據處理成Variable, 如果有GPU, 可以轉成cuda形式
def get_variable(x):
    x = Variable(x)
    return x.cuda() if torch.cuda.is_available() else x

# 對數據進行載入及有相應變換,將Compose看成一種容器,他能對多種數據變換進行組合
# 傳入的參數是一個列表,列表中的元素就是對載入的數據進行的各種變換操作(只有一個顏色通道)
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5,],std=[0.5,])])
# 獲取MNIST訓練集和測試集
data_train=datasets.MNIST(root='data/',transform=transform,train=True,download=True)
data_test=datasets.MNIST(root='data/',transform=transform,train=False)

# 數據裝載
data_loader_train=torch.utils.data.DataLoader(dataset=data_train,batch_size=batch_size,shuffle=True)
data_loader_test = torch.utils.data.DataLoader(dataset =data_test,batch_size = batch_size,shuffle = True)

#數據預覽和圖片顯示
images,labels=next(iter(data_loader_train))
img=torchvision.utils.make_grid(images)
img=img.numpy().transpose(1,2,0)
std=[0.5,0.5,0.5]
mean=[0.5,0.5,0.5]
img=img*std+mean
print([labels[i] for i in range(64)])
plt.imshow(img)
plt.show()

# 構建卷積神經網絡模型
class CNN_Model(torch.nn.Module):
    def __init__(self):
        super(CNN_Model, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(stride=2,kernel_size=2))
        self.conv2=torch.nn.Sequential(
            torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            torch.nn.BatchNorm2d(128),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(stride=2, kernel_size=2))
        self.dense = torch.nn.Sequential(
            torch.nn.Linear(7 * 7 * 128, 1024),
            torch.nn.ReLU(),
            torch.nn.Dropout(p=0.5),
            torch.nn.Linear(1024, 10))
    # 前向傳播
    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        x = x2.view(-1, 7 * 7 * 128)
        x = self.dense(x)
        return x

# 對模型進行訓練和參數優化
cnn_model = CNN_Model()
# 將所有的模型參數移動到GPU上
if torch.cuda.is_available():
    cnn_model = cnn_model.cuda()
loss_func = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(cnn_model.parameters(),lr=learning_rate)
for epoch in range(num_epochs):
    running_loss = 0.0
    running_correct = 0.0
    print("Epoch  {}/{}".format(epoch, num_epochs))
    for data in data_loader_train:
        X_train, y_train = data
        X_train, y_train = get_variable(X_train),get_variable(y_train)
        outputs = cnn_model(X_train)
        _, pred = torch.max(outputs.data, 1)
        optimizer.zero_grad()
        loss = loss_func(outputs, y_train)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        running_correct += torch.sum(pred == y_train.data)
    testing_correct = 0.0
    for data in data_loader_test:
        X_test, y_test = data
        X_test, y_test = get_variable(X_test),get_variable(y_test)
        outputs = cnn_model(X_test)
        _, pred = torch.max(outputs, 1) #返回每一行中最大值的那個元素,且返回其索引
        testing_correct += torch.sum(pred == y_test.data)
        # print(testing_correct)
    print("Loss is :{:.4f},Train Accuracy is:{:.4f}%,Test Accuracy is:{:.4f}%".format(
        running_loss / len(data_train), 100 * running_correct / len(data_train),
        100 * testing_correct / len(data_test)))
# 保存模型
torch.save(cnn_model, 'data/cnn_model.pt')
# 加載模型
cnn_model=torch.load('data/cnn_model.pt')
cnn_model.eval()

運行結果如下:

三、CIFAR10圖像分類

1、CIFAR10數據集介紹

CIFAR-10 是由 Hinton 的學生 Alex Krizhevsky 和 Ilya Sutskever 整理的一個用於識別普適物體的小型數據集。一共包含 10 個類別的 RGB 彩色圖片:飛機( airplane )、汽車( automobile )、鳥類( bird )、貓( cat )、鹿( deer )、狗( dog )、蛙類( frog )、馬( horse )、船( ship )和卡車( truck )。圖片的尺寸爲 32×32 ,數據集中一共有 50000 張訓練圖片和 10000 張測試圖片。與 MNIST 數據集相比, CIFAR-10 具有以下不同點:

  • CIFAR-10 是 3 通道的彩色 RGB 圖像,而 MNIST 是灰度圖像。

  • CIFAR-10 的圖片尺寸爲 32×32, 而 MNIST 的圖片尺寸爲 28×28,比 MNIST 稍大。

  • 相比於手寫字符, CIFAR-10 含有的是現實世界中真實的物體,不僅噪聲很大,而且物體的比例、 特徵都不盡相同,這爲識別帶來很大困難。

2、CIFAR10圖像分類實現

import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
#參數設置
num_epochs = 15
batch_size = 64
learning_rate = 0.001
# 構建CNN模型
class CNNNet(nn.Module):
    def __init__(self):
        super(CNNNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(64, 128, 5)
        self.fc1 = nn.Linear(128* 5 * 5, 1024)
        self.fc2 = nn.Linear(1024, 84)
        self.fc3 = nn.Linear(84, 10)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 128 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
# 圖片顯示
def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()
# torchvision 數據集的輸出是範圍在[0,1]之間的 PILImage,我們將他們轉換成歸一化範圍爲[-1,1]之間的張量Tensors
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 獲取CIFAR10訓練集和測試集
trainset=torchvision.datasets.CIFAR10(root='data/',train=True,download=True,transform=transform)
testset=torchvision.datasets.CIFAR10(root='data/',train=False,download=True,transform=transform)
# CIFAR10訓練集和測試集裝載
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,shuffle=True, num_workers=0)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,shuffle=False, num_workers=0)
# 圖片類別
classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# 圖片顯示
images,labels=next(iter(trainloader))
imshow(torchvision.utils.make_grid(images))

# 定義損失函數和優化器
cnn_model=CNNNet()
criterion=nn.CrossEntropyLoss()
optimizer=optim.SGD(cnn_model.parameters(),lr=learning_rate,momentum=0.9)

# 訓練模型
for epoch in range(num_epochs):
    running_loss=0.00
    running_correct=0.0
    print("Epoch  {}/{}".format(epoch, num_epochs))
    for i,data in enumerate(trainloader,0):
        inputs,labels=data
        optimizer.zero_grad()
        outputs=cnn_model(inputs)
        loss=criterion(outputs,labels)
        loss.backward()
        optimizer.step()
        running_loss+=loss.item()
        _, pred = torch.max(outputs.data, 1)
        running_correct += torch.sum(pred == labels.data)
    print("Loss is :{:.4f},Train Accuracy is:{:.4f}%".format(running_loss / len(trainset), 100 * running_correct / len(trainset)))
# 保存訓練好的模型
torch.save(cnn_model, 'data/cnn_model.pt')

# 加載訓練好的模型
cnn_model=torch.load('data/cnn_model.pt')
cnn_model.eval()
#使用測試集對模型進行評估
correct=0.0
total=0.0
with torch.no_grad():   # 爲了使下面的計算圖不佔用內存
    for data in testloader:
        images, labels = data
        outputs = cnn_model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    print("Test Average accuracy is:{:.4f}%".format(100 * correct / total))

# 求出每個類別的準確率
class_correct=list(0. for i in range(10))
class_total=list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images,labels=data
        outputs=cnn_model(images)
        _,predicted=torch.max(outputs,1)
        c=(predicted==labels).squeeze()
        try:
            for i in range(batch_size):
                label = labels[i]
                class_correct[label] += c[i].item()
                class_total[label] += 1
        except IndexError:
            continue
for i in range(10):
    print('Accuracy of %5s : %4f %%' % (classes[i], 100 * class_correct[i] / class_total[i]))

圖片顯示結果:

模型訓練結果:

測試集平均準確率和每個類別的準確率:

3、在GPU上跑神經網絡

device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
# 遞歸地遍歷所有模塊,並將它們的參數和緩衝器轉換爲CUDA張量
cnn_model.to(device)
# 必須在每一個步驟向GPU發送輸入和目標
inputs,labels=inputs.to(device),labels.to(device)

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章