
PyTorch量化壓縮API:優化深度學習模型的關鍵技術
圖中的結構主要分為兩部分:編碼器和解碼器。編碼器將輸入序列轉換為一組表示,該表示將傳遞給解碼器生成輸出序列。位置編碼(Positional Encoding)是為了在序列中加入位置信息,因為Transformer不具備處理序列位置信息的能力。
PyTorch提供了便捷的API來實現Transformer模型。其核心是將左側的神經網絡封裝為一個TransformerEncoder
類。該類需要兩個關鍵參數:TransformerEncoderLayer
和層數(num_layers
)。
class Transformer(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, num_class,
dim_feedforward=512, num_head=2, num_layers=2, dropout=0.1, max_len=512, activation='relu'):
super(Transformer, self).__init__()
self.embedding_dim = embedding_dim
self.embeddings = nn.Embedding(vocab_size, embedding_dim)
self.position_embedding = PositionalEncoding(embedding_dim, dropout, max_len)
encoder_layer = nn.TransformerEncoderLayer(hidden_dim, num_head, dim_feedforward, dropout, activation)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
self.output = nn.Linear(hidden_dim, num_class)
def forward(self, inputs, lengths):
inputs = torch.transpose(inputs, 0, 1)
hidden_states = self.embeddings(inputs)
hidden_states = self.position_embedding(hidden_states)
attention_mask = length_to_mask(lengths) == False
hidden_states = self.transformer(hidden_states, src_key_padding_mask=attention_mask).transpose(0, 1)
logits = self.output(hidden_states)
log_probs = F.log_softmax(logits, dim=-1)
return log_probs
TransformerEncoder
是一個由多個編碼器層組成的堆棧。每個編碼器層都由TransformerEncoderLayer
類實例化,這個類具體實現了自注意力機制和前饋神經網絡。
class TransformerEncoder(Module):
def __init__(self, encoder_layer, num_layers, norm=None):
super(TransformerEncoder, self).__init__()
self.layers = _get_clones(encoder_layer, num_layers)
self.num_layers = num_layers
self.norm = norm
def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
output = src
for mod in self.layers:
output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
if self.norm is not None:
output = self.norm(output)
return output
該類主要通過克隆傳入的編碼器層并對其進行堆疊,實現了多層編碼器的構建。每層的輸入輸出維度保持一致,這在后續的TransformerEncoderLayer
中可以進一步驗證。
TransformerEncoderLayer
由自注意力機制和前饋網絡組成。其實現基于論文《Attention Is All You Need》,該論文首次提出了Transformer的概念。
class TransformerEncoderLayer(Module):
def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation=F.relu,
layer_norm_eps=1e-5, batch_first=False, norm_first=False,
device=None, dtype=None) -> None:
factory_kwargs = {'device': device, 'dtype': dtype}
super(TransformerEncoderLayer, self).__init__()
self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
**factory_kwargs)
self.linear1 = Linear(d_model, dim_feedforward, **factory_kwargs)
self.dropout = Dropout(dropout)
self.linear2 = Linear(dim_feedforward, d_model, **factory_kwargs)
self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
self.dropout1 = Dropout(dropout)
self.dropout2 = Dropout(dropout)
def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
x = src
x = self.norm1(x + self._sa_block(x, src_mask, src_key_padding_mask))
x = self.norm2(x + self._ff_block(x))
return x
def _sa_block(self, x: Tensor, attn_mask: Optional[Tensor], key_padding_mask: Optional[Tensor]) -> Tensor:
x = self.self_attn(x, x, x, attn_mask=attn_mask, key_padding_mask=key_padding_mask, need_weights=False)[0]
return self.dropout1(x)
def _ff_block(self, x: Tensor) -> Tensor:
x = self.linear2(self.dropout(self.activation(self.linear1(x))))
return self.dropout2(x)
在forward函數中,通過自注意力機制和前饋網絡模塊的運算,輸入的張量經過兩次加法運算后依舊保持了原始的維度,這也證明了TransformerEncoder
中的層堆疊是維度一致的。
在activations.py
文件中,我們可以看到PyTorch對幾種常用激活函數的實現。這些激活函數在深度學習模型中用于引入非線性特性。
class PytorchGELUTanh(nn.Module):
def __init__(self):
super().__init__()
if version.parse(torch.__version__) =1.12.0 is required to use '
'PytorchGELUTanh. Please upgrade torch.'
)
def forward(self, input: Tensor) -> Tensor:
return nn.functional.gelu(input, approximate='tanh')
class NewGELUActivation(nn.Module):
def forward(self, input: Tensor) -> Tensor:
return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0))))
這些實現展示了激活函數的多樣性和PyTorch在實現這些函數時對性能和準確度的權衡。
Transformer是一種用于處理序列數據的深度學習模型,具有很強的并行處理能力。其通過自注意力機制實現了對序列中每個元素的全局依賴建模。
PyTorch通過nn.TransformerEncoder
和nn.TransformerEncoderLayer
等類提供了對Transformer模型的實現。這些類封裝了自注意力機制和前饋網絡的細節,使得用戶可以專注于模型的高層搭建。
用戶可以通過繼承nn.Module
類并實現forward
方法來自定義激活函數。PyTorch提供了多種激活函數的實現作為參考。
由于Transformer模型不具備處理輸入序列中位置信息的能力,因此需要通過位置編碼將位置信息顯式地加入到輸入中,以幫助模型學習序列中的順序關系。
可以通過調整模型的超參數(如層數、隱藏層維度、注意力頭數等),優化數據預處理流程,以及利用GPU加速計算等方法來提高Transformer模型的性能。