如何對loss進行mask

pytorch官方教程中有一個Chatbot教程，就是利用seq2seq和注意力機制實現的，感覺和機器翻譯沒什麼不同啊，如果對話中一句話有下一句，那麼就把這一對句子加入模型進行訓練。其中在訓練階段，損失函式通常需要進行mask操作，因為一個batch中句子的長度通常是不一樣的，一個batch中不足長度的位置需要進行填充（pad）補0，最後生成句子計算loss時需要忽略那些原本是pad的位置的值，即只保留mask中值為1位置的值，忽略值為0位置的值，具體演示如下：

import torch
import torch.nn as nn
import torch.
 
nn.functional as F
import itertools

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

PAD_token = 0

首先是pad函式和建立mask矩陣，矩陣的維度應該和目標一致。

def zeroPadding(l, fillvalue=PAD_token):
    # 輸入：[[1, 1, 1], [2, 2], [3]]
    # 返回：[(1, 2, 3), (1, 2, 0), (1, 0, 0)] 返回已經是轉置後的 [L, B]
    return
 
 list(itertools.zip_longest(*l, fillvalue=fillvalue))


def binaryMatrix(l):
    # 將targets裡非pad部分標記為1，pad部分標記為0
    m = []
    for i, seq in enumerate(l):
        m.append([])
        for token in seq:
            if token == PAD_token:
                m[i].append(0)
            else:
                m[
 
i].append(1)
    return m

假設現在輸入一個batch中有三個句子，我們按照長度從大到小排好序，LSTM或是GRU的輸入和輸出我們需要利用pack_padded_sequence和pad_packed_sequence進行打包和解包，感覺也是在進行mask操作。

inputs = [[1, 2, 3], [4, 5], [6]]   # 輸入句，一個batch，需要按照長度從大到小排好序
inputs_lengths = [3, 2, 1]
targets = [[1, 2], [1, 2, 3], [1]]  # 目標句，這裡的長度是不確定的，mask是針對targets的
inputs_batch = torch.LongTensor(zeroPadding(inputs))
inputs_lengths = torch.LongTensor(inputs_lengths)
targets_batch = torch.LongTensor(zeroPadding(targets))
targets_mask = torch.ByteTensor(binaryMatrix(zeroPadding(targets)))  # 注意這裡是ByteTensor
print(inputs_batch)
print(targets_batch)
print(targets_mask)

列印後結果如下，可見維度統一變成了[L, B]，並且mask和target長得一樣。另外，seq2seq模型處理時for迴圈每次讀取一行，預測下一行的值（即[B, L]時的一列預測下一列）。

tensor([[ 1,  4,  6],
        [ 2,  5,  0],
        [ 3,  0,  0]])
tensor([[ 1,  1,  1],
        [ 2,  2,  0],
        [ 0,  3,  0]])
tensor([[ 1,  1,  1],
        [ 1,  1,  0],
        [ 0,  1,  0]], dtype=torch.uint8)

現在假設我們將inputs輸入模型後，模型讀入sos後預測的第一行為outputs1, 維度為[B, vocab_size]，即每個詞在詞彙表中的概率，模型輸出之前需要softmax。

outputs1 = torch.FloatTensor([[0.2, 0.1, 0.7], [0.3, 0.6, 0.1], [0.4, 0.5, 0.1]])
print(outputs1)

tensor([[ 0.2000,  0.1000,  0.7000],
        [ 0.3000,  0.6000,  0.1000],
        [ 0.4000,  0.5000,  0.1000]])

先看看兩個函式

• torch.gather(input, dim, index, out=None)->Tensor
  沿著某個軸，按照指定維度採集資料，對於3維資料，相當於進行如下操作：

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

比如在這裡，在第1維，選第二個元素。

# 收集每行的第2個元素
temp = torch.gather(outputs1, 1, torch.LongTensor([[1], [1], [1]]))
print(temp)

tensor([[ 0.1000],
        [ 0.6000],
        [ 0.5000]])

• torch.masked_select(input, mask, out=None)->Tensor
  根據mask(ByteTensor)選取對應位置的值，返回一維張量。

例如在這裡我們選取temp大於等於0.5的值。

mask = temp.ge(0.5)  # 大於等於0.5
print(mask)
print(torch.masked_select(temp, temp.ge(0.5)))

tensor([[ 0],
        [ 1],
        [ 1]], dtype=torch.uint8)
tensor([ 0.6000,  0.5000])

然後我們就可以計算loss了，這裡是負對數損失函式，之前模型的輸出要進行softmax。

# 計算一個batch內的平均負對數似然損失，即只考慮mask為1的元素
def maskNLLLoss(inp, target, mask):
    nTotal = mask.sum()
    # 收集目標詞的概率，並取負對數
    crossEntropy = -torch.log(torch.gather(inp, 1, target.view(-1, 1)))
    # 只保留mask中值為1的部分，並求均值
    loss = crossEntropy.masked_select(mask).mean()
    loss = loss.to(DEVICE)
    return loss, nTotal.item()

這裡我們計算第一行的平均損失。

# 計算預測的第一行和targets的第一行的loss
maskNLLLoss(outputs1, targets_batch[0], targets_mask[0])

(tensor(1.1689, device='cuda:0'), 3)

最後進行最後把所有行的loss累加起來變為total_loss.backward()進行反向傳播就可以了。