小知識 | PyTorch 的13個特性

PyTorch在學(xué)術(shù)界和工業(yè)界的應(yīng)用研究中都獲得了很多關(guān)注。它是一個具有很大靈活性的深度學(xué)習(xí)框架，使用了大量的實用工具和函數(shù)來加快工作速度。PyTorch的學(xué)習(xí)曲線并不是那么陡峭，但在其中實現(xiàn)高效和干凈的代碼可能會很棘手。在使用它超過2年之后，以下是我最喜歡的PyTorch功能，我希望我一開始學(xué)習(xí)它就知道。

1. DatasetFolder

當(dāng)學(xué)習(xí)PyTorch時，人們首先要做的事情之一是實現(xiàn)自己的某種Dataset 。這是一個低級錯誤，沒有必要浪費時間寫這樣的東西。通常，數(shù)據(jù)集要么是數(shù)據(jù)列表(或者是numpy數(shù)組)，要么磁盤上的文件。所以，把數(shù)據(jù)在磁盤上組織好，要比寫一個自定義的Dataset來加載某種奇怪的格式更好。

分類器最常見的數(shù)據(jù)格式之一，是有一個帶有子文件夾的目錄，子文件夾表示類，子文件夾中的文件表示樣本，如下所示。

folder/class_0/file1.txt
folder/class_0/file2.txt
folder/class_0/...

folder/class_1/file3.txt
folder/class_1/file4.txt

folder/class_2/file5.txt
folder/class_2/...

有一個內(nèi)置的方式來加載這類數(shù)據(jù)集，不管你的數(shù)據(jù)是圖像，文本文件或其他什么，只要使用'DatasetFolder就可以了。令人驚訝的是，這個類是torchvision包的一部分，而不是核心PyTorch。這個類非常全面，你可以從文件夾中過濾文件，使用自定義代碼加載它們，并動態(tài)轉(zhuǎn)換原始文件。例子:

from torchvision.datasets import DatasetFolder
from pathlib import Path
# I have text files in this folder
ds = DatasetFolder("/Users/marcin/Dev/tmp/my_text_dataset", 
    loader=lambda path: Path(path).read_text(),
    extensions=(".txt",), #only load .txt files
    transform=lambda text: text[:100], # only take first 100 characters
)

# Everything you need is already there
len(ds), ds.classes, ds.class_to_idx
(20, ['novels', 'thrillers'], {'novels': 0, 'thrillers': 1})

如果你在處理圖像，還有一個torchvision.datasets.ImageFolder類，它基于DatasetLoader，它被預(yù)先配置為加載圖像。

2. 盡量少用 .to(device) ，用 zeros_like / ones_like 之類的代替

我讀過很多來自GitHub倉庫的PyTorch代碼。最讓我惱火的是，幾乎在每個repo中都有許多*.to(device)行，它們將數(shù)據(jù)從CPU或GPU轉(zhuǎn)移到其他地方。這樣的語句通常會出現(xiàn)在大量的repos或初學(xué)者教程中。我強烈建議盡可能少地實現(xiàn)這類操作，并依賴內(nèi)置的PyTorch功能自動實現(xiàn)這類操作。到處使用.to(device)通常會導(dǎo)致性能下降，還會出現(xiàn)異常：

Expected object of device type cuda but got device type cpu

顯然，有些情況下你無法回避它，但大多數(shù)情況(如果不是全部)都在這里。其中一種情況是初始化一個全0或全1的張量，這在深度神經(jīng)網(wǎng)絡(luò)計算損失的的時候是經(jīng)常發(fā)生的，模型的輸出已經(jīng)在cuda上了，你需要另外的tensor也是在cuda上，這時，你可以使用*_like操作符：

my_output # on any device, if it's cuda then my_zeros will also be on cuda
my_zeros = torch.zeros_like(my_output_from_model)

在內(nèi)部，PyTorch所做的是調(diào)用以下操作：

my_zeros = torch.zeros(my_output.size(), dtype=my_output.dtype, layout=my_output.layout, device=my_output.device)

所以所有的設(shè)置都是正確的，這樣就減少了代碼中出現(xiàn)錯誤的概率。類似的操作包括：

torch.zeros_like()
torch.ones_like()
torch.rand_like()
torch.randn_like()
torch.randint_like()
torch.empty_like()
torch.full_like()

3. Register Buffer ( nn.Module.register_buffer)

這將是我勸人們不要到處使用 .to(device) 的下一步。有時，你的模型或損失函數(shù)需要有預(yù)先設(shè)置的參數(shù)，并在調(diào)用forward時使用，例如，它可以是一個“權(quán)重”參數(shù)，它可以縮放損失或一些固定張量，它不會改變，但每次都使用。對于這種情況，請使用nn.Module.register_buffer 方法，它告訴PyTorch將傳遞給它的值存儲在模塊中，并將這些值隨模塊一起移動。如果你初始化你的模塊，然后將它移動到GPU，這些值也會自動移動。此外，如果你保存模塊的狀態(tài)，buffers也會被保存！

一旦注冊，這些值就可以在forward函數(shù)中訪問，就像其他模塊的屬性一樣。

from torch import nn
import torch

class ModuleWithCustomValues(nn.Module):
    def __init__(self, weights, alpha):
        super().__init__()
        self.register_buffer("weights", torch.tensor(weights))
        self.register_buffer("alpha", torch.tensor(alpha))
    
    def forward(self, x):
        return x * self.weights + self.alpha

m = ModuleWithCustomValues(
    weights=[1.0, 2.0], alpha=1e-4
)
m(torch.tensor([1.23, 4.56]))
tensor([1.2301, 9.1201])

4. Built-in Identity()

有時候，當(dāng)你使用遷移學(xué)習(xí)時，你需要用1:1的映射替換一些層，可以用nn.Module來實現(xiàn)這個目的，只返回輸入值。PyTorch內(nèi)置了這個類。

例子，你想要在分類層之前從一個預(yù)訓(xùn)練過的ResNet50獲取圖像表示。以下是如何做到這一點：

from torchvision.models import resnet50
model = resnet50(pretrained=True)
model.fc = nn.Identity()
last_layer_output = model(torch.rand((1, 3, 224, 224)))
last_layer_output.shape
torch.Size([1, 2048])

5. Pairwise distances: torch.cdist

下次當(dāng)你遇到計算兩個張量之間的歐幾里得距離(或者一般來說:p范數(shù))的問題時，請記住torch.cdist。它確實做到了這一點，并且在使用歐幾里得距離時還自動使用矩陣乘法，從而提高了性能。

points1 = torch.tensor([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]])
points2 = torch.tensor([[0.0, 0.0], [-1.0, -1.0], [-2.0, -2.0], [-3.0, -3.0]]) # batches don't have to be equal
torch.cdist(points1, points2, p=2.0)
tensor([[0.0000, 1.4142, 2.8284, 4.2426],
        [1.4142, 2.8284, 4.2426, 5.6569],
        [2.8284, 4.2426, 5.6569, 7.0711]])

沒有矩陣乘法或有矩陣乘法的性能，在我的機器上使用mm時，速度快了2倍以上。

%%timeit
points1 = torch.rand((512, 2))
points2 = torch.rand((512, 2))
torch.cdist(points1, points2, p=2.0, compute_mode="donot_use_mm_for_euclid_dist")

867μs±142μs per loop (mean±std. dev. of 7 run, 1000 loop each)

%%timeit
points1 = torch.rand((512, 2))
points2 = torch.rand((512, 2))
torch.cdist(points1, points2, p=2.0)

417μs±52.9μs per loop (mean±std. dev. of 7 run, 1000 loop each)

6. Cosine similarity: F.cosine_similarity

與上一點相同，計算歐幾里得距離并不總是你需要的東西。當(dāng)處理向量時，通常余弦相似度是選擇的度量。PyTorch也有一個內(nèi)置的余弦相似度實現(xiàn)。

import torch.nn.functional as F
vector1 = torch.tensor([0.0, 1.0])
vector2 = torch.tensor([0.05, 1.0])
print(F.cosine_similarity(vector1, vector2, dim=0))
vector3 = torch.tensor([0.0, -1.0])
print(F.cosine_similarity(vector1, vector3, dim=0))
tensor(0.9988)
tensor(-1.)

PyTorch中批量計算余弦距離

import torch.nn.functional as F
batch_of_vectors = torch.rand((4, 64))
similarity_matrix = F.cosine_similarity(batch_of_vectors.unsqueeze(1), batch_of_vectors.unsqueeze(0), dim=2)
similarity_matrix
tensor([[1.0000, 0.6922, 0.6480, 0.6789],
        [0.6922, 1.0000, 0.7143, 0.7172],
        [0.6480, 0.7143, 1.0000, 0.7312],
        [0.6789, 0.7172, 0.7312, 1.0000]])

7. 歸一化向量: F.normalize

最后一點仍然與向量和距離有松散的聯(lián)系，那就是歸一化：通常是通過改變向量的大小來提高計算的穩(wěn)定性。最常用的歸一化是L2，可以在PyTorch中按如下方式應(yīng)用:

vector = torch.tensor([99.0, -512.0, 123.0, 0.1, 6.66])
normalized_vector = F.normalize(vector, p=2.0, dim=0)
normalized_vector
tensor([ 1.8476e-01, -9.5552e-01,  2.2955e-01,  1.8662e-04,  1.2429e-02])

在PyTorch中執(zhí)行歸一化的舊方法是：

vector = torch.tensor([99.0, -512.0, 123.0, 0.1, 6.66])
normalized_vector = vector / torch.norm(vector, p=2.0)
normalized_vector
tensor([ 1.8476e-01, -9.5552e-01,  2.2955e-01,  1.8662e-04,  1.2429e-02])

在PyTorch中批量進行L2歸一化

batch_of_vectors = torch.rand((4, 64))
normalized_batch_of_vectors = F.normalize(batch_of_vectors, p=2.0, dim=1)
normalized_batch_of_vectors.shape, torch.norm(normalized_batch_of_vectors, dim=1) # all vectors will have length of 1.0
(torch.Size([4, 64]), tensor([1.0000, 1.0000, 1.0000, 1.0000]))

8. 線性層 + 分塊技巧 (torch.chunk)

這是我最近發(fā)現(xiàn)的一個有創(chuàng)意的技巧。假設(shè)你想把你的輸入映射到N個不同的線性投影中。你可以通過創(chuàng)建N個nn.Linear來做到這一點。或者你也可以創(chuàng)建一個單一的線性層，做一個向前傳遞，然后將輸出分成N塊。這種方法通常會帶來更高的性能，所以這是一個值得記住的技巧。

d = 1024
batch = torch.rand((8, d))
layers = nn.Linear(d, 128, bias=False), nn.Linear(d, 128, bias=False), nn.Linear(d, 128, bias=False)
one_layer = nn.Linear(d, 128 * 3, bias=False)
%%timeit
o1 = layers[0](batch)
o2 = layers[1](batch)
o3 = layers[2](batch)

289 μs ± 30.8 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%%timeit
o1, o2, o3 = torch.chunk(one_layer(batch), 3, dim=1)

202 μs ± 8.09 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

9. Masked select (torch.masked_select)

有時你只需要對輸入張量的一部分進行計算。給你一個例子：你想計算的損失只在滿足某些條件的張量上。為了做到這一點，你可以使用torch.masked_select，注意，當(dāng)需要梯度時也可以使用這個操作。

data = torch.rand((3, 3)).requires_grad_()
print(data)
mask = data > data.mean()
print(mask)
torch.masked_select(data, mask)
tensor([[0.0582, 0.7170, 0.7713],
        [0.9458, 0.2597, 0.6711],
        [0.2828, 0.2232, 0.1981]], requires_grad=True)
tensor([[False,  True,  True],
        [ True, False,  True],
        [False, False, False]])
tensor([0.7170, 0.7713, 0.9458, 0.6711], grad_fn=<MaskedSelectBackward>)

直接在tensor上應(yīng)用mask

類似的行為可以通過使用mask作為輸入張量的 “indexer”來實現(xiàn)。

data[mask]
tensor([0.7170, 0.7713, 0.9458, 0.6711], grad_fn=<IndexBackward>)

有時，一個理想的解決方案是用0填充mask中所有的False值，可以這樣做:

data * mask
tensor([[0.0000, 0.7170, 0.7713],
        [0.9458, 0.0000, 0.6711],
        [0.0000, 0.0000, 0.0000]], grad_fn=<MulBackward0>)

10. 使用 torch.where來對tensors加條件

當(dāng)你想把兩個張量結(jié)合在一個條件下這個函數(shù)很有用，如果條件是真，那么從第一個張量中取元素，如果條件是假，從第二個張量中取元素。

x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], requires_grad=True)
y = -x
condition_or_mask = x <= 3.0
torch.where(condition_or_mask, x, y)
tensor([ 1.,  2.,  3., -4., -5.], grad_fn=<SWhereBackward>)

11. 在給定的位置給張量填入值(Tensor.scatter)

這個函數(shù)的用例如下，你想用給定位置下另一個張量的值填充一個張量。一維張量更容易理解，所以我將先展示它，然后繼續(xù)更高級的例子。

data = torch.tensor([1, 2, 3, 4, 5])
index = torch.tensor([0, 1])
values = torch.tensor([-1, -2, -3, -4, -5])
data.scatter(0, index, values)
tensor([-1, -2,  3,  4,  5])

上面的例子很簡單，但是現(xiàn)在看看如果將index改為index = torch.tensor([0, 1, 4])會發(fā)生什么：

data = torch.tensor([1, 2, 3, 4, 5])
index = torch.tensor([0, 1, 4])
values = torch.tensor([-1, -2, -3, -4, -5])
data.scatter(0, index, values)
tensor([-1, -2,  3,  4, -3])

為什么最后一個值是-3，這是反直覺的，對吧？這是PyTorch scatter函數(shù)的中心思想。index變量表示data張量的第i個值應(yīng)該放在values張量的哪個位置。我希望下面的簡單python版的這個操作能讓你更明白：

data_orig = torch.tensor([1, 2, 3, 4, 5])
index = torch.tensor([0, 1, 4])
values = torch.tensor([-1, -2, -3, -4, -5])
scattered = data_orig.scatter(0, index, values)

data = data_orig.clone()
for idx_in_values, where_to_put_the_value in enumerate(index):
    what_value_to_put = values[idx_in_values]
    data[where_to_put_the_value] = what_value_to_put
data, scattered
(tensor([-1, -2,  3,  4, -3]), tensor([-1, -2,  3,  4, -3]))

2D數(shù)據(jù)的PyTorch scatter例子

始終記住，index的形狀與values的形狀相關(guān)，而index中的值對應(yīng)于data中的位置。

data = torch.zeros((4, 4)).float()
index = torch.tensor([
    [0, 1],
    [2, 3],
    [0, 3],
    [1, 2]
])
values = torch.arange(1, 9).float().view(4, 2)
values, data.scatter(1, index, values)
(tensor([[1., 2.],
        [3., 4.],
        [5., 6.],
        [7., 8.]]),
tensor([[1., 2., 0., 0.],
        [0., 0., 3., 4.],
        [5., 0., 0., 6.],
        [0., 7., 8., 0.]]))

12. 在網(wǎng)絡(luò)中進行圖像插值 (F.interpolate)

當(dāng)我學(xué)習(xí)PyTorch時，我驚訝地發(fā)現(xiàn)，實際上可以在前向傳遞中調(diào)整圖像(或任何中間張量)，并保持梯度流。這種方法在使用CNN和GANs時特別有用。

# image from https://commons.wikimedia.org/wiki/File:A_female_British_Shorthair_at_the_age_of_20_months.jpg
img = Image.open("./cat.jpg")
img

to_pil_image(
    F.interpolate(to_tensor(img).unsqueeze(0),  # batch of size 1
                  mode="bilinear", 
                  scale_factor=2.0, 
                  align_corners=False).squeeze(0) # remove batch dimension
)

看看梯度流是如何保存的：

F.interpolate(to_tensor(img).unsqueeze(0).requires_grad_(),
                  mode="bicubic", 
                  scale_factor=2.0, 
                  align_corners=False)
tensor([[[[0.9216, 0.9216, 0.9216,  ..., 0.8361, 0.8272, 0.8219],
    [0.9214, 0.9214, 0.9214,  ..., 0.8361, 0.8272, 0.8219],
    [0.9212, 0.9212, 0.9212,  ..., 0.8361, 0.8272, 0.8219],
    ...,
    [0.9098, 0.9098, 0.9098,  ..., 0.3592, 0.3486, 0.3421],
    [0.9098, 0.9098, 0.9098,  ..., 0.3566, 0.3463, 0.3400],
    [0.9098, 0.9098, 0.9098,  ..., 0.3550, 0.3449, 0.3387]],

    [[0.6627, 0.6627, 0.6627,  ..., 0.5380, 0.5292, 0.5238],
    [0.6626, 0.6626, 0.6626,  ..., 0.5380, 0.5292, 0.5238],
    [0.6623, 0.6623, 0.6623,  ..., 0.5380, 0.5292, 0.5238],
    ...,
    [0.6196, 0.6196, 0.6196,  ..., 0.3631, 0.3525, 0.3461],
    [0.6196, 0.6196, 0.6196,  ..., 0.3605, 0.3502, 0.3439],
    [0.6196, 0.6196, 0.6196,  ..., 0.3589, 0.3488, 0.3426]],

    [[0.4353, 0.4353, 0.4353,  ..., 0.1913, 0.1835, 0.1787],
    [0.4352, 0.4352, 0.4352,  ..., 0.1913, 0.1835, 0.1787],
    [0.4349, 0.4349, 0.4349,  ..., 0.1913, 0.1835, 0.1787],
    ...,
    [0.3333, 0.3333, 0.3333,  ..., 0.3827, 0.3721, 0.3657],
    [0.3333, 0.3333, 0.3333,  ..., 0.3801, 0.3698, 0.3635],
    [0.3333, 0.3333, 0.3333,  ..., 0.3785, 0.3684, 0.3622]]]],
grad_fn=<UpsampleBicubic2DBackward1>)

13. 將圖像做成網(wǎng)格 (torchvision.utils.make_grid)

當(dāng)使用PyTorch和torchvision時，不需要使用matplotlib或一些外部庫來復(fù)制粘貼代碼來顯示圖像網(wǎng)格。只要使用torchvision.utils.make_grid就行了。

from torchvision.utils import make_grid
from torchvision.transforms.functional import to_tensor, to_pil_image
from PIL import Image
img = Image.open("./cat.jpg")
to_pil_image(
    make_grid(
        [to_tensor(i) for i in [img, img, img]],
         nrow=2, # number of images in single row
         padding=5 # "frame" size
     )
)