【動手學深度學習】第五章筆記：層與塊、引數管理、自定義層、讀寫檔案、GPU

bringlu發表於2023-04-27
深度學習筆記GPU

為了更好的閱讀體驗，請點選這裡

由於本章內容比較少且以後很顯然會經常回來翻，因此會寫得比較詳細。

5.1 層和塊

事實證明，研究討論“比單個層大”但“比整個模型小”的元件更有價值。例如，在計算機視覺中廣泛流行的ResNet-152 架構就有數百層，這些層是由層組（groups of layers）的重複模式組成。

為了實現這些複雜的網路，我們引入了神經網路塊的概念。塊（block）可以描述單個層、由多個層組成的元件或整個模型本身。使用塊進行抽象的一個好處是可以將一些塊組合成更大的元件。透過定義程式碼來按需生成任意複雜度的塊，我們可以透過簡潔的程式碼實現複雜的神經網路。

從程式設計的角度來看，塊由類（class）表示。它的任何子類都必須定義一個將其輸入轉換為輸出的前向傳播函式，並且必須儲存任何必需的引數。注意，有些塊不需要任何引數。最後，為了計算梯度，塊必須具有反向傳播函式。在定義我們自己的塊時，由於自動微分提供了一些後端實現，我們只需要考慮前向傳播函式和必需的引數。

之後原書中舉的例子為例項化一個包含兩個線性層的多層感知機。該程式碼中，透過例項化 nn.Sequential 來構建模型，層的執行順序是作為引數傳遞的。簡而言之，nn.Sequential 定義了一種特殊的 Module，即在 PyTorch 中表示一個塊的類，它維護了一個由 Module 組成的有序列表。注意，兩個全連線層都是 Linear 類的例項，Linear 類本身就是 Module 的子類。另外，到目前為止，我們一直在透過 net(X) 呼叫我們的模型來獲得模型的輸出。這實際上是 net.__call__(X) 的簡寫。

5.1.1 自定義塊

實現自定義塊之前，簡要總結一下每個塊必須提供的基本功能。

  1. 將輸入資料作為其前向傳播函式的引數。
  2. 透過前向傳播函式來生成輸出。請注意，輸出的形狀可能與輸入的形狀不同。例如，我們上面模型中的第一個全連線的層接收一個20維的輸入，但是返回一個維度為256的輸出。
  3. 計算其輸出關於輸入的梯度，可透過其反向傳播函式進行訪問。通常這是自動發生的。
  4. 儲存和訪問前向傳播計算所需的引數。
  5. 根據需要初始化模型引數。

在下面的程式碼片段中，我們從零開始編寫一個塊。它包含一個多層感知機，其具有 \(256\) 個隱藏單元的隱藏層和一個 \(10\) 維輸出層。注意，下面的 MLP 類繼承了表示塊的類。我們的實現只需要提供我們自己的建構函式（Python中的 __init__ 函式）和前向傳播函式。

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20, 256)
        self.out = nn.Linear(256, 10)
    
    def forward(self, X):
        return self.out(F.relu(self.hidden(X)))

注意一些關鍵細節：首先，我們定製的 __init__ 函式透過 super().__init__() 呼叫父類的 __init__ 函式，省去了重複編寫模版程式碼的痛苦。然後，我們例項化兩個全連線層，分別為 self.hidden 和 self.out。注意，除非我們實現一個新的運算子，否則我們不必擔心反向傳播函式或引數初始化，系統將自動生成這些。

塊的一個主要優點是它的多功能性。我們可以子類化塊以建立層（如全連線層的類）、整個模型（如上面的MLP類）或具有中等複雜度的各種元件。

5.1.2 順序塊

構建簡化的 MySequential，只需要定義兩個關鍵函式：

  1. 一種將塊逐個追加到列表中的函式；
  2. 一種前向傳播函式，用於將輸入按追加塊的順序傳遞給塊組成的“鏈條”。

下面的 MySequential 類提供了與預設 Sequential 類相同的功能。

class MySequential(nn.Module):
    def __init__(self, *args):
        super().__init__()
        for idx, module in enumerate(args):
            # 這裡，module 是 Module 子類的一個例項。我們把它儲存在 'Module' 類的成員
            # 變數 _modules 中。_module 的型別是 OrderedDict
            self._modules[str(idx)] = module

    def forward(self, X):
        # OrderedDict 保證了按照成員新增的順序遍歷它們
        for block in self._modules.values():
            X = block(X)
        return X

__init__ 函式將每個模組逐個新增到有序字典 _modules 中。讀者可能會好奇為什麼每個 Module 都有一個 _modules 屬性？以及為什麼我們使用它而不是自己定義一個Python列表？簡而言之，_modules 的主要優點是：在模組的引數初始化過程中，系統知道在 _modules 字典中查詢需要初始化引數的子塊。

5.1.3 在前向傳播函式中執行程式碼

當需要更強的靈活性時，我們需要定義自己的塊。例如，可能希望在前向傳播函式中執行Python的控制流。此外，可能希望執行任意的數學運算，而不是簡單地依賴預定義的神經網路層。

那麼，就可以在前向傳播的函式中實現複雜的程式碼。

練習題

（1）如果將 MySequential 中儲存塊的方式更改為 Python 列表，會出現什麼樣的問題？

class MySequential(nn.Module):
    def __init__(self, *args):
        super().__init__()
        self.modules_list = []
        for idx, module in enumerate(args):
            self.modules_list.append(module)
        print(self.modules_list)
    
    def forward(self, X):
        for block in self.modules_list:
            X = block(X)
        return X
    
net = MySequential(nn.Linear(20, 256), nn.ReLU(), nn.Linear(256, 10))
net(X)

接下來如果呼叫 net.parameters() 迭代器來遍歷引數或者用 net.state_dict() 來檢視狀態字典，你會發現什麼也不會輸出。原因在於 parameter 型別的引數只能從 _modules 中以及其他顯示定義在表層的 nn.Module 類及子類獲得，即使你把 list 換成另一個 OrderedDict 也並不好用。現在沒辦法自動獲取了。

除此之外，由於無法自動獲取 parameter 型別的引數，因此初始化很難做。

（2）實現一個塊，它以兩個塊為引數，例如 net1 和 net2，並返回前向傳播中兩個網路的串聯輸出。這也被稱為平行塊。

class ParallelBlock(nn.Module):
    def __init__(self, net1, net2):
        super().__init__()
        self.net1 = net1
        self.net2 = net2
    
    def forward(self, X):
        return self.net2(self.net1(X))
    
net = ParallelBlock(nn.Linear(16, 20), nn.Linear(20, 10))
print(net)
for param in net.parameters():
    print(param)

（3）假設我們想要連線同一網路的多個例項。實現一個函式，該函式生成同一個塊的多個例項，並在此基礎上構建更大的網路。

一般而言 Sequential 就足夠完成這個任務：

class multilayer(nn.Module):
    def __init__(self, num):
        super().__init__()
        layer_list = []
        for i in range(num):
            layer_list.append(nn.Linear(20, 10))
        self.ln = nn.Sequential(*layer_list)
    
    def forward(self, X):
        return self.ln(X)

multilayer(
  (ln): Sequential(
    (0): Linear(in_features=20, out_features=10, bias=True)
    (1): Linear(in_features=20, out_features=10, bias=True)
    (2): Linear(in_features=20, out_features=10, bias=True)
    (3): Linear(in_features=20, out_features=10, bias=True)
    (4): Linear(in_features=20, out_features=10, bias=True)
  )
)

當然，也可以使用 nn.ModuleList：

class multilayer(nn.Module):
    def __init__(self, num):
        super().__init__()
        layer_list = []
        for i in range(num):
            layer_list.append(nn.Linear(20, 10))
        self.ln = nn.ModuleList(layer_list)
    
    def forward(self, X):
        return self.ln(X)

multilayer(
  (ln): ModuleList(
    (0): Linear(in_features=20, out_features=10, bias=True)
    (1): Linear(in_features=20, out_features=10, bias=True)
    (2): Linear(in_features=20, out_features=10, bias=True)
    (3): Linear(in_features=20, out_features=10, bias=True)
    (4): Linear(in_features=20, out_features=10, bias=True)
  )
)

5.2 引數管理

有時我們希望提取引數，以便在其他環境中複用它們，將模型儲存下來，以便它可以在其他軟體中執行，或者為了獲得科學的理解而進行檢查。

本節，我們將介紹以下內容：

  • 訪問引數，用於除錯、診斷和視覺化；
  • 引數初始化；
  • 在不同模型元件間共享引數。

假定此時有一個單隱藏層的多層感知機

import torch
from torch import nn

net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 1))
X = torch.rand(size = (2, 4))
net(X)

tensor([[-0.5471], [-0.5554]], grad_fn=<AddmmBackward0>)

5.2.1 引數訪問

同時，對於 Sequential 中，可以使用索引來訪問模型的任意層，除此之外，可以使用 .state_dict() 來檢查引數。比如，第二個全連線層的呼叫方法為 net[2].state_dict()。

OrderedDict([('weight', tensor([[-0.2183, -0.2935, -0.2471,  0.3105, -0.0285, -0.0140, -0.1047, -0.0894]])), ('bias', tensor([-0.0456]))])

1. 目標引數

parameter 是複合的類，包含值、梯度和額外資訊。這就是我們需要顯式引數值的原因。除了值之外，我們還可以訪問每個引數的梯度。

print(type(net[2].bias))
print(net[2].bias)
print(net[2].bias.data)

<class 'torch.nn.parameter.Parameter'>
Parameter containing:
tensor([0.2615], requires_grad=True)
tensor([0.2615])

2. 一次性訪問所有引數

當我們需要對所有引數執行操作時，逐個訪問它們可能會很麻煩。當我們處理更復雜的塊（例如，巢狀塊）時，情況可能會變得特別複雜，因為我們需要遞迴整個樹來提取每個子塊的引數。下面，我們將透過演示來比較訪問第一個全連線層的引數和訪問所有層。

module.named_parameters 返回一個所有 module 引數的迭代器，返回引數名字和引數。

print(*[(name, param.shape) for name, param in net[0].named_parameters()])
print(*[(name, param.shape) for name, param in net.named_parameters()])

('weight', torch.Size([8, 4])) ('bias', torch.Size([8]))
('0.weight', torch.Size([8, 4])) ('0.bias', torch.Size([8])) ('2.weight', torch.Size([1, 8])) ('2.bias', torch.Size([1]))

也有另一種訪問網路引數的方式：

net.state_dict()['2.bias'].data

tensor([0.2615])

3. 從巢狀塊收集引數

def block1():
    return nn.Sequential(nn.Linear(4, 8), nn.ReLU(),
                        nn.Linear(8, 4), nn.ReLU())

def block2():
    net = nn.Sequential()
    for i in range(4):
        net.add_module(f'block {i}', block1())
    return net

rgnet = nn.Sequential(block2(), nn.Linear(4, 1))
rgnet(X)

tensor([[0.2608],
        [0.2611]], grad_fn=<AddmmBackward0>)

輸出一下看看

print(rgnet)

Sequential(
  (0): Sequential(
    (block 0): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
    (block 1): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
    (block 2): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
    (block 3): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
  )
  (1): Linear(in_features=4, out_features=1, bias=True)
)

由於是巢狀了三層 Sequential 因此可以使用索引來訪問層。

rgnet[0][1][0].bias.data

tensor([-0.0647,  0.1259, -0.3926, -0.3025, -0.1323,  0.3075,  0.4889,  0.1187])

5.2.2 引數初始化

深度學習框架提供預設隨機初始化，也允許我們建立自定義初始化方法，滿足我們透過其他規則實現初始化權重。

預設情況下，PyTorch 會根據一個範圍均勻地初始化權重和偏置矩陣，這個範圍是根據輸入和輸出維度計算出的。PyTorch 的 nn.init 模組提供了多種預置初始化方法。

1. 內建初始化

首先呼叫內建的初始化器。下面的程式碼將所有權重引數初始化為標準差為 \(0.01\) 的高斯隨機變數，且將偏置引數設定為 \(0\)。

def init_normal(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight, mean=0, std=0.01)
        nn.init.zeros_(m.bias)
net.apply(init_normal)
net[0].weight.data[0], net[0].bias.data[0]

(tensor([-0.0261,  0.0005,  0.0169,  0.0050]), tensor(0.))

還可以將所有引數初始化為給定的常量，如初始化為 \(1\)。

def init_constant(m):
    if type(m) == nn.Linear:
        nn.init.constant_(m.weight, 1)
        nn.init.zeros_(m.bias)
net.apply(init_constant)
net[0].weight.data[0], net[0].bias[0]

(tensor([1., 1., 1., 1.]), tensor(0., grad_fn=<SelectBackward0>))

我們還可以對某些塊應用不同的初始化方法。例如，下面我們使用 Xavier 初始化方法初始化第一個神經網路層，然後將第三個神經網路層初始化為常量值 \(42\)。

def init_xavier(m):
    if type(m) == nn.Linear:
        nn.init.xavier_uniform_(m.weight)
def init_42(m):
    if type(m) == nn.Linear:
        nn.init.constant_(m.weight, 42)

net[0].apply(init_xavier)
net[2].apply(init_42)
print(net[0].weight.data[0])
print(net[2].weight.data)

tensor([ 0.3676,  0.3810,  0.5257, -0.0244])
tensor([[42., 42., 42., 42., 42., 42., 42., 42.]])

2. 自定義初始化

有時，深度學習框架沒有提供我們需要的初始化方法。在下面的例子中，使用以下的分佈為任意權重引數 \(w\) 定義初始化方法：

\[w \sim \begin{cases} U(5, 10), &\text{可能性} \frac{1}{4} \\ 0, &\text{可能性}\frac{1}{2} \\ U(-10, -5), &\text{可能性} \frac{1}{4} \end{cases} \]

同樣，實現了一個 my_init 函式來應用到 net。

def my_init(m):
    if type(m) == nn.Linear:
        print("Init", *[(name, param.shape) for name, param in m.named_parameters()][0])
        nn.init.uniform_(m.weight, -10, 10)
        m.weight.data *= m.weight.data.abs() >= 5
        
net.apply(my_init)
net[0].weight[:2]

Init weight torch.Size([8, 4])
Init weight torch.Size([1, 8])
tensor([[-7.2929, -0.0000, -0.0000, -5.2074],
        [ 9.1947, -8.8687,  0.0000,  0.0000]], grad_fn=<SliceBackward0>)

注意，始終可以直接設定引數。

net[0].weight.data[:] += 1
net[0].weight.data[0, 0] = 42
net[0].weight.data[0]

tensor([42.0000,  1.0000,  1.0000, -4.2074])

5.2.3 引數繫結

有時我們希望在多個層間共享引數：我們可以定義一個稠密層，然後使用它的引數來設定另一個層的引數。

# 我們需要給共享層一個名稱，以便可以引用它的引數
shared = nn.Linear(8, 8)
net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(),
                    shared, nn.ReLU(),
                    shared, nn.ReLU(),
                    nn.Linear(8, 1))
net(X)
# 檢查引數是否相同
print(net[2].weight.data[0] == net[4].weight.data[0])
net[2].weight.data[0, 0] = 100
# 確保它們實際上是同一個物件，而不只是有相同的值
print(net[2].weight.data[0] == net[4].weight.data[0])

tensor([True, True, True, True, True, True, True, True])
tensor([True, True, True, True, True, True, True, True])

這個例子表明第三個和第五個神經網路層的引數是繫結的。它們不僅值相等，而且由相同的張量表示。因此，如果我們改變其中一個引數，另一個引數也會改變。這裡有一個問題：當引數繫結時，梯度會發生什麼情況？答案是由於模型引數包含梯度，因此在反向傳播期間第二個隱藏層（即第三個神經網路層）和第三個隱藏層（即第五個神經網路層）的梯度會加在一起。

練習題

（1）使用之前沒寫的 NestMLP (FancyMLP) 模型訪問各個層的引數。

class NestMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(20, 64), nn.ReLU(),
                                 nn.Linear(64, 32), nn.ReLU())
        self.linear = nn.Linear(32, 16)

    def forward(self, X):
        return self.linear(self.net(X))

net = NestMLP()
for name, param in net.named_parameters():
    print(name, param.shape)

net.0.weight torch.Size([64, 20])
net.0.bias torch.Size([64])
net.2.weight torch.Size([32, 64])
net.2.bias torch.Size([32])
linear.weight torch.Size([16, 32])
linear.bias torch.Size([16])

（2）檢視初始化模組文件以瞭解不同的初始化方法。

官方文件連結

（3）構建包含共享引數層的多層感知機並對其進行訓練。在訓練過程中，觀察模型各層的引數和梯度。

舉個簡單的例子，\(z=wy, y=wx\)，不妨假設此時複製了兩個與 \(w\) 相同的值 \(w_1, w_2\)。那麼在反向傳播中 \(\frac{\mathrm{d} z}{\mathrm{d} w} = \frac{\mathrm{d}z}{\mathrm{d} w_1} + \frac{\mathrm{d} z}{\mathrm{d} y} \frac{\mathrm{d} y}{\mathrm{d} w_2} = y + wx = 2wx\)，因此會是多倍梯度加和。

（4）為什麼共享引數是個好方式？

可以減少引數，空間佔用更小。但是正確性有待商榷。

5.3 延後初始化

延後初始化（defers initialization），即直到資料第一次透過模型傳遞時，框架才會動態地推斷出每個層的大小。

在以後，當使用卷積神經網路時，由於輸入維度（即影像的解析度）將影響每個後續層的維數，有了該技術將更加方便。現在我們在編寫程式碼時無須知道維度是什麼就可以設定引數，這種能力可以大大簡化定義和修改模型的任務。

延後初始化中只有第一層需要延遲初始化，但是框架仍是按順序初始化的。等到知道了所有的引數形狀，框架就可以初始化引數。

書上沒有關於延後初始化的程式碼，原因在於 PyTorch 中的延後初始化層 nn.LazyLinear() 仍然還是一個開發中的 feature。所以這一節在 PyTorch 版的書裡有什麼存在的必要嗎？

5.4 自定義層

本節將展示如何構建自定義層。

5.4.1 不帶引數的層

首先，構造一個沒有任何引數的自定義層。下面的 CenteredLayer 類要從其輸入中減去均值。要構建它，我們只需繼承基礎層類並實現前向傳播功能。

class CenteredLayer(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, X):
        return X - X.mean()

5.4.2 帶引數的層

下面繼續定義具有引數的層， 這些引數可以透過訓練進行調整。可以使用內建函式來建立引數，這些函式提供一些基本的管理功能。比如管理訪問、初始化、共享、儲存和載入模型引數。這樣做的好處之一是：我們不需要為每個自定義層編寫自定義的序列化程式。

下面實現自定義版本的全連線層：

class MyLinear(nn.Module):
    def __init__(self, in_units, units):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(in_units, units))
        self.bias = nn.Parameter(torch.randn(units,))
    def forward(self, X):
        linear = torch.matmul(X, self.weight.data) + self.bias.data
        return F.relu(linear)

練習題

（1）設計一個接受輸入並計算張量降維的層，它返回 \(y_k = \sum_{i,j} W_{ijk} x_i x_j\)

最好使用 transpose() 或者是 permute() 把 \(W_{ijk}\) 轉換一個維度，變成 \(W_{kij}\)。這樣就可以寫成如下的形式了：

\[y_k = \boldsymbol{x}^T \boldsymbol{W}_k \boldsymbol{x} \]

class testlayer1(nn.Module):
    def __init__(self, in_units, units):
        super().__init__()
        self.W = nn.Parameter(torch.randn(units, in_units, in_units))
    def forward(self, x):
        h1 = torch.matmul(x, self.W.data)
        h2 = torch.matmul(h1, x)
        return h2
    
net = testlayer1(4, 2)
a = torch.rand(4)
print(a, net(a))
# 驗證一下第一個對不對
print(torch.matmul(a, torch.matmul(net.W[0], a)))

tensor([0.2971, 0.8508, 0.0615, 0.5073]) tensor([-0.5827, -1.1151])
tensor(-0.5827, grad_fn=<DotBackward0>)

第二題看不懂 QWQ

5.5 讀寫檔案

5.5.1 載入和儲存張量

本節內容為如何載入和儲存權重向量和整個模型。

  • torch.save(obj, f) 儲存張量 obj 到 f 位置。
  • torch.load(f) 讀取 f 位置的檔案。

書中給出了儲存與讀取張量、張量列表、張量字典的示例。

5.5.2 載入和儲存模型引數

深度學習框架提供了內建函式來儲存和載入整個網路。需要注意的一個重要細節是，這將儲存模型的引數而不是儲存整個模型。例如，如果有一個 \(3\) 層多層感知機，則需要單獨指定架構。因為模型本身可以包含任意程式碼，所以模型本身難以序列化。因此，為了恢復模型，需要用程式碼生成架構，然後從磁碟載入引數。從多層感知機開始：

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20, 256)
        self.output = nn.Linear(256, 10)
    def forward(self, x):
        return self.output(F.relu(self.hidden(x)))

net = MLP()
X = torch.randn(size=(2, 20))
Y = net(X)

接下來，將模型的引數 net.state_dict() 儲存在一個 mlp.params 的檔案中。

torch.save(net.state_dict(), 'mlp.params')

為了恢復模型，我們例項化了原始多層感知機模型的一個備份。這裡不需要隨機初始化模型引數，而是直接讀取檔案中儲存的引數。

clone = MLP()
clone.load_state_dict(torch.load('mlp.params'))

這樣即完成了模型的儲存和載入。

練習題

（1）即使不需要將經過訓練的模型部署到不同的裝置上，儲存模型引數還有什麼實際的好處？

可以讓其他人複用模型，做重複實驗。

（2）假設我們只想複用網路的一部分，以將其合併到不同的網路架構中。比如想在一個新的網路中使用之前網路的前兩層，該怎麼做？

這裡僅使用上文中多層感知機的第一層作為例子。

old_net_state_dict = torch.load('mlp.params')
clone2 = MLP()
# 假設此處預處理剩下層已經完成
clone2.hidden.weight.data = old_net_state_dict["hidden.weight"]
clone2.hidden.bias.data = old_net_state_dict["hidden.bias"]

或者直接從這個基於 OrderedDict 的 state_dict 裡面拿引數就行。

（3）如何同時儲存網路架構和引數？需要對架構加上什麼限制？

直接 torch.save(net) 即可。但是這個網路架構不包括 forward 函式。

5.6 GPU

可以使用 nvidia-smi 命令來檢視顯示卡資訊。

我用的 Kaggle 平臺的 T4 2 張，可以完成本節的程式碼任務。

!nvidia-smi

Thu Apr 27 09:27:16 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.161.03   Driver Version: 470.161.03   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   36C    P8     9W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla T4            Off  | 00000000:00:05.0 Off |                    0 |
| N/A   34C    P8    10W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

5.6.1 計算裝置

在 PyTorch 中，CPU 和 GPU 可以用 torch.device('cpu') 和 torch.device('cuda') 表示。應該注意的是，cpu 裝置意味著所有物理 CPU 和記憶體，這意味著 PyTorch 的計算將嘗試使用所有 CPU 核心。然而，gpu 裝置只代表一個卡和相應的視訊記憶體。如果有多個 GPU，我們使用 torch.device(f'cuda:{i}') 來表示第 \(i\) 塊 GPU（\(i\) 從 \(0\) 開始）。另外，cuda:0 和 cuda 是等價的。

import torch
from torch import nn

torch.device('cpu'), torch.device('cuda'), torch.device('cuda:1')

(device(type='cpu'), device(type='cuda'), device(type='cuda', index=1))

還可以查詢可用的 GPU 的數量。

torch.cuda.device_count()

2

原書中定義了兩個方便的函式，這兩個函式允許在不存在所需 GPU 的情況下執行程式碼。

  • try_gpu(i) 嘗試使用 \(i\) 號 GPU，如果存在返回 torch.device(f'cuda:{i}')，如果不存在返回 torch.device('cpu')。預設引數為 i=0。
  • try_all_gpus() 嘗試使用所有 GPU，如果存在 GPU 返回所有 GPU 的列表，如果不存在返回 [torch.device('cpu')]。

5.6.2 張量與 GPU

預設情況下，張量是在 CPU 上建立的。需要注意的是，無論何時我們要對多個項進行操作，它們都必須在同一個裝置上。

1. 儲存在 GPU 上

有幾種方法可以在 GPU 上儲存張量。例如，我們可以在建立張量時指定儲存裝置。接下來，我們在第一個 gpu 上建立張量變數 X。在 GPU 上建立的張量只消耗這個 GPU 的視訊記憶體。我們可以使用 nvidia-smi 命令檢視視訊記憶體使用情況。 一般來說，我們需要確保不建立超過 GPU 視訊記憶體限制的資料。

X = torch.ones(2, 3, device = try_gpu())
X

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')

假設還存在另一個 GPU，那麼在另一個 GPU 上建立隨機張量。

Y = torch.rand(2, 3, device = try_gpu(1))
Y

tensor([[0.4099, 0.3582, 0.8877],
        [0.7732, 0.8459, 0.1519]], device='cuda:1')

2. 複製

如果要計算 \(\sf X + Y\)，那麼需要將它們弄到同一個裝置上，然後才能執行運算操作。例如，下面的程式碼是將 \(\sf X\) 複製到第二個 GPU，然後執行加法運算。

Z = X.cuda(1)
print(X)
print(Z)

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')
tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:1')

當然，也可以使用 .to() 來執行復制：

Z = X.to(torch.device('cuda:1'))
Z

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:1')

相加：

Y + Z

tensor([[1.4099, 1.3582, 1.8877],
        [1.7732, 1.8459, 1.1519]], device='cuda:1')

假設變數 \(\sf Z\) 已經存在於第二個 GPU 上。如果我們還是呼叫 Z.cuda(1) 會發生什麼？它將返回 \(\sf Z\)，而不會複製並分配新記憶體。

Z.cuda(1) is Z

True

注意呼叫 Z.to(torch.device("cuda:1")) is Z 也同樣返回 True。

所以這個 .to() 和 .cuda() 有啥區別啊

5.6.3 神經網路與 GPU

類似地，可以神經網路模型可以指定裝置。下面的程式碼將模型引數放在 GPU 上。

net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=try_gpu())
net(X)

tensor([[-0.3980],
        [-0.3980]], device='cuda:0', grad_fn=<AddmmBackward0>)

練習題

只做第（4）題。

（4）測量同時在兩個 GPU 上執行兩個矩陣乘法與在一個 GPU 上按順序執行兩個矩陣乘法所需的時間。提示：應該看到近乎線性的縮放。

同時在兩個 GPU 上執行矩陣乘法：

a = torch.rand(1000, 1000).to(try_gpu(0))
b = torch.rand(1000, 1000).to(try_gpu(0))
c = torch.rand(1000, 1000).to(try_gpu(1))
d = torch.rand(1000, 1000).to(try_gpu(1))
begintime = time.time()
for i in range(1000):
    e = torch.matmul(a, b)
    f = torch.matmul(c, d)
print(time.time() - begintime)

0.34023451805114746

在一個 GPU 上按順序執行兩個矩陣乘法所需的時間：

a = torch.rand(1000, 1000).to(try_gpu(0))
b = torch.rand(1000, 1000).to(try_gpu(0))
c = torch.rand(1000, 1000).to(try_gpu(0))
d = torch.rand(1000, 1000).to(try_gpu(0))
begintime = time.time()
for i in range(1000):
    e = torch.matmul(a, b)
    f = torch.matmul(c, d)
print(time.time() - begintime)

0.8642914295196533

差不多是兩倍的差距。

相關文章