PyTorch實戰(3)——PyTorch vs. TensorFlow詳解
- 0. 前言
- 1. 張量
- 2. PyTorch 模塊
- 2.1 torch.nn
- 2.2 torch.optim
- 2.3 torch.utils.data
- 3. 使用 PyTorch 訓練神經網絡
- 小結
- 系列鏈接
0. 前言
PyTorch
是一個基于 Torch
庫的 Python
機器學習庫,廣泛用于深度學習的科研和應用開發,主要由 Meta
開發。PyTorch
是另一個知名深度學習庫 TensorFlow
(由 Google
開發)的有力競爭者,最初,這兩者的主要區別在于,PyTorch
基于即時執行 (eager execution
),而 TensorFlow 1.x
基于圖計算的延遲執行 (deferred execution
),但現在 TensorFlow2.x
也提供了即時執行模式。
即時執行基本上是一種命令式編程模式,在這種模式下,數學操作會立即計算。而延遲執行模式會將所有操作存儲在計算圖中,不立即計算,直到構建完成才對整個圖進行評估。即時執行的優勢在于其直觀的流程、易于調試以及更少的輔助代碼。
PyTorch
通過類似 NumPy
的語法/接口,提供了張量計算能力,并能利用 GPU
實現加速計算。張量是計算單元,類似于 NumPy
數組,不同之處在于張量可以在 GPU
上使用,以加速計算。
憑借加速計算能力和創建動態計算圖,PyTorch
提供了一個完整的深度學習框架。除此之外,它具有真正的 Python
風格,使 PyTorch
用戶能夠充分利用 Python
的所有特性,包括豐富的 Python
數據科學生態系統。
在本節中,我們將深入探討張量的概念及其在 PyTorch
中的實現方式,以及張量具有的屬性。我們還將了解一些常用的 PyTorch
模塊,這些模塊擴展了數據加載、模型構建以及訓練過程中優化算法的功能。我們將這些 PyTorch API
與 TensorFlow
的對應功能進行對比,以了解兩者在底層實現上的差異。
1. 張量
張量 (Tensor
) 的概念類似于 NumPy
數組。張量是一個 nnn 維數組,我們可以對其執行數學函數、通過 GPU
加速計算,還可以跟蹤計算圖和梯度,這些功能對深度學習至關重要。為了在 GPU
上運行張量,只需將張量轉換為特定的數據類型即可。
(1) 使用 PyTorch
實例化一個張量:
points = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
(2) 獲取第一個元素:
points[0]
(4) 查看張量的形狀:
points.shape
(5) 使用 TensorFlow
聲明一個張量:
points_tf = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0])
訪問第一個元素或獲取張量形狀的命令與 PyTorch
相同。
(6) 在 PyTorch
中,張量是對存儲在連續內存塊中的一維數組的視圖,這些數組稱為存儲實例。每個 PyTorch
張量都有一個 untyped_storage()
屬性,可以調用它來輸出張量的底層存儲實例:
points = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
points.untyped_storage()
輸出結果如下所示:
TensorFlow
張量沒有 untyped_storage()
屬性。PyTorch
張量是存儲實例的視圖,張量使用以下信息來實現該視圖:
- 大小 (
Size
) - 存儲 (
Storage
) - 偏移量 (
Offset
) - 步長 (
Stride
)
這些信息的含義如下:
size
類似于NumPy
中的shape
屬性,表示每個維度上的元素數量:
這些數字的乘積等于底層存儲實例的長度(本例中為points.size()
6
):torch.Size([3, 2])
storage
:底層存儲實例offset
:張量的第一個元素在存儲數組中的索引stride
:表示在每個維度上移動一個元素所需的步長
在 TensorFlow
中,張量的形狀可以通過使用 shape
屬性獲取:
points_tf = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
points_tf.shape
輸出結果如下所示:
TensorShape([3, 2])
我們已經了解了 PyTorch
張量的 untyped_storage()
屬性,接下來我們看看偏移量 (offset
):
points.storage_offset()
輸出結果如下所示:
0
offset
表示張量的第一個元素在存儲數組中的索引。由于輸出為 0
,這意味著張量的第一個元素是存儲數組的第一個元素。使用以下代碼進行驗證:
points[1].storage_offset()
輸出結果如下所示:
2
因為 points[1]
是 [3.0, 4.0]
,而存儲數組是 [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]
,可以看到張量中的第一個元素 [3.0, 4.0]
,在存儲數組中的索引位置為 2
。
接下來,查看 stride
屬性:
points.stride()
輸出結果如下所示:
(2, 1)
可以看到,stride
屬性表示在每個維度上訪問下一個元素需要跳過的元素數量。因此,在以上例子中,沿著第一個維度,若要訪問第一個元素之后的元素(即 1.0
),我們需要跳過 2
個元素(即 1.0
和 2.0
),才能訪問下一個元素,即 3.0
。類似地,沿著第二個維度,我們需要跳過 1
個元素,才能訪問 1.0
之后的元素,即 2.0
。因此,使用這些屬性,張量可以從一個連續的一維存儲數組中推導出來。TensorFlow
張量并沒有 stride
或 storage_offset
屬性。
(7) 張量中包含的數據是數值類型的。具體來說,PyTorch
提供了以下幾種數據類型供張量使用:
torch.float32
或torch.float
—32
位浮點數torch.float64
或torch.double
—64
位雙精度浮點數torch.float16
或torch.half
—16
位半精度浮點數torch.int8
— 有符號8
位整數torch.uint8
— 無符號8
位整數torch.int16
或torch.short
— 有符號16
位整數torch.int32
或torch.int
— 有符號32
位整數torch.int64
或torch.long
— 有符號64
位整數
TensorFlow 也提供了類似的數據類型。
(8) 在 PyTorch
中,可以使用 dtype
屬性為張量指定特定數據類型:
points = torch.tensor([[1.0,2.0],[3.0,4.0]], dtype=torch.float16)
在 TensorFlow
中,可以通過以下等效代碼實現:
points_tf = tf.constant([[1.0,2.0],[3.0,4.0]], dtype=tf.float16)
(9) 除了數據類型,PyTorch
中的張量可以指定存儲設備:
points = torch.tensor([[1.0,2.0],[3.0,4.0]], dtype=torch.float16, device='cuda')
我們也可以在目標設備上創建張量的副本:
points_2 = points.to(device='cuda')
從以上示例可以看出,我們可以將張量分配到 CPU
(使用 device='cpu'
),如果不指定設備,默認情況下會分配到 CPU
,也可以將張量分配給 GPU
(使用 device='cuda'
)。在 TensorFlow
中,可以使用以下方式分配設備:
with tf.device('/gpu:0'):points_tf = tf.constant([[1.0,2.0],[3.0,4.0]], dtype=tf.float16)
PyTorch
支持 NVIDIA
(CUDA
) 和 AMD GPU
。
當張量被放置在 GPU
上時,計算速度會顯著加快,并且由于 PyTorch
中的張量 API
在 CPU
和 GPU
張量間基本一致,因此可以很方便的在設備之間移動張量、執行計算并移回。
如果有多個相同類型的設備,比如多個 GPU
,我們可以通過設備索引精確指定張量放置的設備:
points_3 = points.to(device='cuda:0')
接下來,我們將介紹一些用于構建深度學習模型的重要 PyTorch
模塊。
2. PyTorch 模塊
PyTorch
庫不僅提供了類似 NumPy
的計算功能,還提供了一系列模塊,幫助開發者快速設計、訓練和測試深度學習模型。
2.1 torch.nn
在構建神經網絡架構時,網絡的基本組成要素包括層數、每層的神經元數量,以及哪些神經元是可學習的等。PyTorch
的 nn
模塊允許用戶通過定義這些高層次特性快速實例化神經網絡架構,而不需要手動指定所有的細節。如果不使用 nn
模塊,需要使用以下方式進行單層神經網絡初始化:
import math
weights = torch.randn(256, 4) / math.sqrt(256)
weights. requires_grad_()
bias = torch.zeros(4, requires_grad=True)
而使用 nn
模塊,nn.Linear(256, 4)
就可以實現相同的功能。在 TensorFlow
中,可以使用以下方式實現:
tf.keras.layers.Dense(256, input_shape=(4, ), activation=None)
在 torch.nn
模塊中,有一個 torch.nn.functional
子模塊,包含了 torch.nn
模塊中的所有函數,包括損失函數、激活函數,以及用于以函數式方式創建神經網絡的函數(即將每一層表示為前一層輸出的函數),例如池化、卷積和線性函數。使用 torch.nn.functional
模塊定義損失函數:
import torch.nn.functional as F
loss_func = F.cross_entropy
loss = loss_func(model(X), y)
其中,X
是輸入,y
是目標輸出,model
是神經網絡模型。在 TensorFlow
中,上述代碼可以寫成:
import tensorflow as tf
loss_func = tf.keras.losses.SparseCategoricalCrossentropy()
loss = loss_func(y, model(X))
2.2 torch.optim
在訓練神經網絡時,我們通過反向傳播調整網絡的權重或參數,這一過程稱為優化。optim
模塊包含了與訓練深度學習模型時運行各種優化算法相關的工具和功能。
使用 torch.optim
模塊在訓練過程中定義優化器:
opt = optim.SGD(model.parameters(), lr=lr)
如果,我們手動編寫優化步驟:
with torch.no_grad():for param in model.parameters():param -= param.grad * lrmodel.zero_grad()
使用優化器可以簡潔地寫成如下形式:
opt.step()
opt.zero_grad()
TensorFlow
不需要顯式地編寫梯度更新和清除步驟,使用優化器代碼如下:
opt = tf.keras.optimizers.SGD(learning_rate=lr)
model.compile(optimizer=opt, loss=loss)
2.3 torch.utils.data
在 utils.data
模塊下,PyTorch
提供了 Dataset
和 DataLoader
類,這些類因其抽象且靈活的實現而非常實用,這些類提供了直觀的方式來迭代數據和其他操作。通過使用這些類,我們可以確保高性能的張量計算,并實現可靠的數據輸入/輸出。可以通過以下方式使用 torch.utils.data.DataLoader
:
from torch.utils.data import TensorDataset, DataLoader
train_dataset = TensorDataset(x_train, y_train)
train_dataloader = DataLoader(train_dataset, batch_sise)
使用這種方式,我們就不需要手動遍歷數據批次:
for i in range((n-1)//bs + 1):x_batch = x_train[start_i:end_i]y_batch = y_train[start_i:end_i]pred = model(x_batch)
我們可以簡單地寫成:
for x_batch, y_batch in train_dataloader:pred = model(x_batch)
torch.utils.data
類似于 TensorFlow
中的 tf.data.Dataset
。在 TensorFlow
中,遍歷數據批次的代碼如下:
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataloader = train_dataset.batch(bs)
for x_batch, y_batch in train_dataloader:pred = model(x_batch)
我們已經了解了 PyTorch
庫,并了解了 PyTorch
的張量。接下來,我們將學習如何使用 PyTorch
訓練神經網絡。
3. 使用 PyTorch 訓練神經網絡
在本節中,我們將使用 MNIST
數據集,該數據集包含手寫郵政編碼數字( 0
到 9
)的圖像及其對應的標簽。MNIST
數據集包含 60000
個訓練樣本和 10000
個測試樣本,每個樣本都是一張 28x28
像素的灰度圖像。PyTorch
在 Dataset
模塊中提供了 MNIST
數據集。我們將使用 PyTorch
在 MNIST
數據集上訓練一個深度學習多類分類器,并測試訓練后的模型在測試樣本上的表現。
(1) 首先,導入所需庫:
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transformsimport matplotlib.pyplot as plt
(2) 接下來,定義模型架構,如下圖所示:
該模型由卷積層、Dropout
層以及線性(全連接)層組成,這些層都可以通過 torch.nn
模塊實現:
class ConvNet(nn.Module):def __init__(self):super(ConvNet, self).__init__()self.cn1 = nn.Conv2d(1, 16, 3, 1)self.cn2 = nn.Conv2d(16, 32, 3, 1)self.dp1 = nn.Dropout(0.10)self.dp2 = nn.Dropout(0.25)self.fc1 = nn.Linear(4608, 64) # 4608 is basically 12 X 12 X 32self.fc2 = nn.Linear(64, 10)def forward(self, x):x = self.cn1(x)x = F.relu(x)x = self.cn2(x)x = F.relu(x)x = F.max_pool2d(x, 2)x = self.dp1(x)x = torch.flatten(x, 1)x = self.fc1(x)x = F.relu(x)x = self.dp2(x)x = self.fc2(x)op = F.log_softmax(x, dim=1)return op
__init__
函數定義了模型的核心架構,forward
函數執行網絡的前向傳播,它包括了每層的激活函數以及池化或 Dropout
操作,函數返回模型的最終輸出(即預測值),其維度與目標輸出(真實標簽)相同。
第一個卷積層的輸入通道為 1
(灰度圖像),輸出通道為 16
,卷積核大小為 3
,步幅為 1
。輸入通道為 1
是因為輸入圖像是灰度圖,選擇 3x3
的卷積核基于以下原因:
- 卷積核的大小通常是奇數,以便輸入圖像的像素圍繞中心像素對稱分布
1x1
的卷積核太小,無法捕捉相鄰像素的信息3x3
是計算機視覺問題中最常用的卷積核大小之一,因為它能夠捕捉局部視覺特征
不選擇 5
、7
或者 27
是由于當卷積核的大小過大時,比如 27x27
,在 28x28
的圖像上卷積,會得到非常粗略的特征。然而,圖像中的最重要視覺特征通常是局部的(在較小的空間鄰域內),因此使用一個小卷積核逐步查看鄰近像素來提取視覺模式更為合理。3x3
是 CNN
中解決計算機視覺問題時最常用的卷積核大小之一。
需要注意的是,我們使用兩個連續的卷積層,兩個卷積層的卷積核大小都是 3x3
。從空間覆蓋的角度來看,這相當于使用一個 5x5
的卷積核進行一次卷積。然而,通常更傾向于使用多個小卷積核的層,因為這樣可以構建更深的網絡,從而學習到更復雜的特征,同時,由于卷積核較小,參數也較少。通過在多個層中使用多個小卷積核,會得到專門檢測不同特征的卷積核——例如有些用于檢測邊緣,有些用于檢測圓形,有些用于檢測紅色等。
第一個卷積層輸入的是單通道數據,輸出 16
個通道。這意味著該層正在嘗試從輸入圖像中提取 16
種不同類型的信息。每個輸出通道稱為特征圖,每個特征圖都有一個專門的卷積核來提取其對應的特征。
第二個卷積層中將通道數從 16
增加到 32
,旨在從圖像中提取更多種類的特征。卷積層輸出通道設計通常遵循先增加后減小的原則。
本節中,我們將步幅 (stride
) 設置為 1
,因為卷積核大小僅為 3
。如果步幅值過大,卷積核會跳過圖像中的許多像素,這不利于特征提取。如果卷積核大小是 100
,那么可能會考慮將步幅設置為 10。步幅越大,卷積操作的次數越少,卷積核的視野 (field of view
) 也會越小。
以上代碼也可以使用 torch.nn.Sequential API
來編寫:
model = nn.Sequential(nn.Conv2d(1, 16, 3, 1),nn.ReLU(),nn.Conv2d(16, 32, 3,1 ),nn.ReLU(),nn.MaxPool2d(2),nn.Dropout(0.1),nn.Flatten(),nn.Linear(4068, 64),nn.ReLU(),nn.Dropout(0.25),nn.Linear(64, 10),nn.LogSoftmax(dim=1)
)
通常推薦通過單獨的 __init__
和 forward
方法來初始化模型,以便在模型層并非順序執行(例如并行或跳躍連接)時,能夠更靈活地定義模型功能。使用 Sequential
的代碼與使用 TensorFlow
非常相似:
import tensorflow as tf
model = tf.keras.Sequential([tf.keras.layers.Conv2D(16, 3, activation='relu', input_shape=(28, 28, 1)),tf.keras.layers.Conv2D(32, 3, activation='relu'),tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),tf.keras.layers.Dropout(0.1),tf.keras.layers.Flatten(),tf.keras.layers.Dense(64, activation='relu'),tf.keras.layers.Dropout(0.25),tf.keras.layers.Dense(10, activation='softmax')
])
而使用 __init__
和 forward
方法的代碼在 TensorFlow
中如下:
class ConvNet(tf.keras.Model):def __init__(self):super(ConvNet, self).__init__()self.cn1 = tf.keras.layers.Conv2D(16, 3, activation='relu', input_shape=(28, 28, 1))self.cn2 = tf.keras.layers.Conv2D(32, 3, activation='relu')self.dp1 = tf.keras.layers.Dropout(0.10)self.dp2 = tf.keras.layers.Dropout(0.25)self.flatten = tf.keras.layers.Flatten()self.fc1 = tf.keras.layers.Dense(64, activation='relu')self.fc2 = tf.keras.layers.Dense(10, activation='softmax')def call(self, x):x = self.cn1(x)x = self.cn2(x)x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)x = self.dp1(x)x = self.flatten(x)x = self.fc1(x)x = self.dp2(x)x = self.fc2(x)return x
在 TensorFlow
中,使用 call
方法代替 forward
,其余部分與 PyTorch
代碼類似。
(3) 接下來,定義訓練過程,即實際的反向傳播步驟。可以看到,torch.optim
模塊極大地簡化了代碼:
def train(model, device, train_dataloader, optim, epoch):model.train()for b_i, (X, y) in enumerate(train_dataloader):X, y = X.to(device), y.to(device)optim.zero_grad()pred_prob = model(X)loss = F.nll_loss(pred_prob, y) # nll is the negative likelihood lossloss.backward()optim.step()if b_i % 10 == 0:print('epoch: {} [{}/{} ({:.0f}%)]\t training loss: {:.6f}'.format(epoch, b_i * len(X), len(train_dataloader.dataset),100. * b_i / len(train_dataloader), loss.item()))
train()
函數以批次遍歷數據集,將數據復制到指定設備上,通過神經網絡模型進行前向傳播,計算模型預測值與真實標簽之間的損失,使用優化器調整模型權重,并每 10
個批次打印一次訓練日志。整個過程執行一次稱為一個 epoch
,即完整遍歷一次數據集對于 TensorFlow
,我們可以直接以高級方式直接運行訓練。PyTorch
中詳細的訓練過程定義使我們能夠更靈活地控制訓練過程,而不是用一行高級代碼完成訓練。
(4) 與訓練過程類似,編寫一個測試過程,用于評估模型在測試集上的表現:
def test(model, device, test_dataloader):model.eval()loss = 0success = 0with torch.no_grad():for X, y in test_dataloader:X, y = X.to(device), y.to(device)pred_prob = model(X)loss += F.nll_loss(pred_prob, y, reduction='sum').item() # loss summed across the batchpred = pred_prob.argmax(dim=1, keepdim=True) # us argmax to get the most likely predictionsuccess += pred.eq(y.view_as(pred)).sum().item()loss /= len(test_dataloader.dataset)print('\nTest dataset: Overall Loss: {:.4f}, Overall Accuracy: {}/{} ({:.0f}%)\n'.format(loss, success, len(test_dataloader.dataset),100. * success / len(test_dataloader.dataset)))
test()
函數的大部分內容與 train()
函數類似。唯一的區別是,計算出的模型預測與真實標簽之間的損失不會用來調整模型權重,而是用于計算整個測試批次的總體測試誤差。
(5) 接下來,加載數據集。得益于 PyTorch
的 DataLoader
模塊,我們可以方便的設置數據集加載機制:
train_dataloader = torch.utils.data.DataLoader(datasets.MNIST('./data', train=True, download=True,transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1302,), (0.3069,))])), # train_X.mean()/256. and train_X.std()/256.batch_size=32, shuffle=True)test_dataloader = torch.utils.data.DataLoader(datasets.MNIST('./data', train=False, transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1302,), (0.3069,)) ])),batch_size=500, shuffle=False)
可以看到,我們將批大小 batch_size
設置為 32
。通常,批大小的選擇需要權衡:太小的批大小會導致訓練速度變慢,因為需要頻繁計算梯度,且梯度噪聲較大;太大的批大小也會因等待梯度計算時間過長而減慢訓練速度。通常不建議等待太長時間才進行一次梯度更新,更頻繁但精度較低的梯度更新最終會引導模型學習到更好的參數。
對于訓練集和測試集,我們指定了數據集保存的存儲位置,并且設置了批大小,批大小決定了每次訓練和測試運行中數據實例的數量。此外,我們還隨機打亂訓練數據實例,以確保數據樣本在各個批次中均勻分布。
最后,將數據集歸一化,使其符合具有指定均值和標準差的正態分布。如果我們從零開始訓練模型,那么均值和標準差來自于訓練數據集,如果我們是從一個預訓練模型進行遷移學習,那么均值和標準差值將來自于預訓練模型的原始訓練數據集。
在 TensorFlow
中,我們可以使用 tf.keras.datasets
來加載 MNIST
數據,并使用 tf.data.Dataset
模塊從數據集中創建訓練數據批次:
# Load the MNIST dataset.
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data() # Normalize pixel values between 0 and 1
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0# Add a channels dimension (required for CNN)
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]# Create a dataloader for training.
train_dataloader = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataloader = train_dataloader.shuffle(10000)
train_dataloader = train_dataloader.batch(32)# Create a dataloader for testing.
test_dataloader = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_dataloader = test_dataloader.batch(500)
(6) 定義優化器和設備,使用它們運行模型訓練:
device = torch.device("cuda")model = ConvNet().to(device=device)
optimizer = optim.Adadelta(model.parameters(), lr=0.5)
在本節中,使用 Adadelta
作為優化器,學習率設置為 0.5
。我們在介紹優化器時提到,如果我們處理的是稀疏數據,選用 Adadelta
可以得到不錯的結果。MNIST
數據集就是一個稀疏數據的例子,因為并非圖像中的所有像素都具有信息量。但我們也可以嘗試其他優化器,如 Adam
,來解決這個問題,觀察不同優化器對訓練過程和模型性能的影響。在 TensorFlow
中,可以使用以下等效代碼實例化并編譯模型:
model = ConvNet()
optimizer = tf.keras.optimizers.experimental.Adadelta(learning_rate=0.5)
model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])
(7) 然后開始實際的模型訓練過程,訓練多個 epoch
,并在每個訓練 epoch
結束時測試模型:
for epoch in range(1, 3):train(model, device, train_dataloader, optimizer, epoch)test(model, device, test_dataloader)
訓練過程輸出結果如下所示:
TensorFlow
中的訓練循環等效代碼如下:
model.fit(train_dataloader, epochs=2, validation_data=test_dataloader)
(8) 模型訓練完成后,我們可以手動檢查模型在樣本圖像上的推理結果是否正確:
test_samples = enumerate(test_dataloader)
b_i, (sample_data, sample_targets) = next(test_samples)plt.imshow(sample_data[0][0], cmap='gray', interpolation='none')
plt.show()
輸出結果如下所示:
TensorFlow
中的等效代碼基本相同,唯一不同的是使用 sample_data[0]
而不是 sample_data[0][0]
:
test_samples = enumerate(test_dataloader)
b_i, (sample_data, sample_targets) = next(test_samples)
plt.imshow(sample_data[0], cmap='gray', interpolation='none')
plt.show()
將圖像輸入訓練后的模型,運行模型推理,并比較預測結果與真實標簽:
print(f"Model prediction is : {model(sample_data.to(device=device)).data.max(1)[1][0]}")
print(f"Ground truth is : {sample_targets[0]}")
需要注意的是,對于預測,首先使用 max()
函數在 axis=1
軸上計算概率最大的類別。max()
函數會輸出兩個列表——sample_data
中每個樣本的類別概率列表和每個樣本的類別標簽列表。因此,我們使用索引 [1]
選擇第二個列表(即類別標簽列表),并通過索引 [0]
進一步選擇第一個類別標簽,以查看 sample_data
中的第一個樣本。輸出結果如下所示:
Model prediction is : 7
Ground truth is : 7
可以看到,得到了正確的預測結果。神經網絡的前向傳播通過 model()
完成,會得到類別概率。因此,我們使用 max()
函數輸出最大概率對應的類別。在 TensorFlow
中,可以使用以下代碼獲取預測結果:
print(f"Model prediction is : {tf.math.argmax(model(sample_data)[0])}")
print(f"Ground truth is : {sample_targets[0]}")
小結
在本節中,我們比較了 PyTorch
和 TensorFlow
兩大深度學習庫,并在模型訓練的不同階段(模型初始化、數據加載、訓練循環和模型評估)分析了 PyTorch
與 TensorFlow
的 API
差異,最后,作為實踐分別使用 PyTorch
和 TensorFlow
從零開始訓練了一個深度學習模型。
系列鏈接
PyTorch實戰(1)——深度學習概述
PyTorch實戰(2)——使用PyTorch構建神經網絡
PyTorch實戰(4)——卷積神經網絡(Convolutional Neural Network,CNN)