Shortcuts

开始使用嵌套张量

创建于:2022年8月2日 | 最后更新:2024年5月3日 | 最后验证:2024年11月5日

嵌套张量扩展了常规密集张量的形状,允许表示不规则大小的数据。

  • 对于一个常规的张量,每个维度都是规则的并且有一个大小

  • 对于一个嵌套张量,并非所有维度都有规则的大小;其中一些是不规则的

嵌套张量是表示各个领域中序列数据的自然解决方案:

  • 在自然语言处理中,句子的长度可以是可变的,因此一批句子会形成一个嵌套张量

  • 在计算机视觉中,图像可以有不同的形状,因此一批图像形成一个嵌套张量

在本教程中,我们将演示嵌套张量的基本用法,并通过一个实际例子说明它们在处理不同长度序列数据时的实用性。特别是,它们在构建能够高效处理不规则序列输入的变压器时非常宝贵。下面,我们展示了一个使用嵌套张量实现的多头注意力机制,结合torch.compile的使用,其性能优于在填充张量上的简单操作。

嵌套张量目前是一个原型功能,可能会发生变化。

import numpy as np
import timeit
import torch
import torch.nn.functional as F

from torch import nn

torch.manual_seed(1)
np.random.seed(1)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

嵌套张量初始化

从Python前端,可以从张量列表中创建一个嵌套张量。 我们将nt[i]表示为嵌套张量的第i个张量组件。

nt = torch.nested.nested_tensor([torch.arange(12).reshape(
    2, 6), torch.arange(18).reshape(3, 6)], dtype=torch.float, device=device)
print(f"{nt=}")

通过将每个底层张量填充到相同的形状,可以将嵌套张量转换为常规张量。

padded_out_tensor = torch.nested.to_padded_tensor(nt, padding=0.0)
print(f"{padded_out_tensor=}")

所有张量都有一个属性,用于确定它们是否是嵌套的;

print(f"nt is nested: {nt.is_nested}")
print(f"padded_out_tensor is nested: {padded_out_tensor.is_nested}")

通常从形状不规则的张量批次中构建嵌套张量。 即假设维度0是批次维度。 索引维度0会返回第一个基础张量组件。

print("First underlying tensor component:", nt[0], sep='\n')
print("last column of 2nd underlying tensor component:", nt[1, :, -1], sep='\n')

# When indexing a nestedtensor's 0th dimension, the result is a regular tensor.
print(f"First underlying tensor component is nested: {nt[0].is_nested}")

一个重要注意事项是,目前尚未支持在维度0上进行切片。 这意味着目前无法构建一个结合底层张量组件的视图。

嵌套张量操作

由于每个操作都必须为嵌套张量明确实现,目前嵌套张量的操作覆盖范围比常规张量要窄。目前,仅涵盖了索引、dropout、softmax、转置、重塑、线性、bmm等基本操作。然而,覆盖范围正在扩大。如果您需要某些操作,请提交一个issue以帮助我们优先考虑覆盖范围。

reshape

reshape操作用于改变张量的形状。 关于常规张量的完整语义可以在这里找到。 对于常规张量,当指定新形状时, 单个维度可以为-1,在这种情况下,它从剩余的维度和元素数量中推断出来。

嵌套张量的语义类似,除了-1不再推断。 相反,它继承了旧的大小(这里nt[0]为2,nt[1]为3)。 -1是指定锯齿维度的唯一合法大小。

nt_reshaped = nt.reshape(2, -1, 2, 3)
print(f"{nt_reshaped=}")

转置

转置操作用于交换张量的两个维度。 其完整的语义可以在这里找到。 请注意,对于嵌套张量,维度0是特殊的; 它被假定为批次维度, 因此不支持涉及嵌套张量维度0的转置。

nt_transposed = nt_reshaped.transpose(1, 2)
print(f"{nt_transposed=}")

其他

其他操作与常规张量具有相同的语义。 在嵌套张量上应用操作等同于 对基础张量组件应用操作, 结果也是一个嵌套张量。

nt_mm = torch.nested.nested_tensor([torch.randn((2, 3, 4)), torch.randn((2, 3, 5))], device=device)
nt3 = torch.matmul(nt_transposed, nt_mm)
print(f"Result of Matmul:\n {nt3}")

nt4 = F.dropout(nt3, 0.1)
print(f"Result of Dropout:\n {nt4}")

nt5 = F.softmax(nt4, -1)
print(f"Result of Softmax:\n {nt5}")

为什么使用嵌套张量

当数据是顺序的时,通常每个样本的长度都不同。 例如,在一批句子中,每个句子的单词数量不同。 处理不同序列的常见技术是手动填充每个数据张量 以达到相同的形状,从而形成一个批次。 例如,我们有两个长度不同的句子和一个词汇表 为了将其表示为单个张量,我们用0填充到批次中的最大长度。

sentences = [["goodbye", "padding"],
             ["embrace", "nested", "tensor"]]
vocabulary = {"goodbye": 1.0, "padding": 2.0,
              "embrace": 3.0, "nested": 4.0, "tensor": 5.0}
padded_sentences = torch.tensor([[1.0, 2.0, 0.0],
                                 [3.0, 4.0, 5.0]])
nested_sentences = torch.nested.nested_tensor([torch.tensor([1.0, 2.0]),
                                               torch.tensor([3.0, 4.0, 5.0])])
print(f"{padded_sentences=}")
print(f"{nested_sentences=}")

这种将一批数据填充到其最大长度的技术并不是最优的。填充的数据在计算中并不需要,并且通过分配比必要更大的张量来浪费内存。此外,并非所有操作在应用于填充数据时都具有相同的语义。对于矩阵乘法,为了忽略填充的条目,需要用0填充,而对于softmax,则必须用-inf填充以忽略特定条目。嵌套张量的主要目标是促进在不规则数据上使用标准的PyTorch张量UX进行操作,从而消除对低效和复杂的填充和掩码的需求。

padded_sentences_for_softmax = torch.tensor([[1.0, 2.0, float("-inf")],
                                             [3.0, 4.0, 5.0]])
print(F.softmax(padded_sentences_for_softmax, -1))
print(F.softmax(nested_sentences, -1))

让我们来看一个实际的例子:在Transformers中使用的多头注意力组件。 我们可以以这样一种方式实现它,使其可以在填充或嵌套张量上操作。

class MultiHeadAttention(nn.Module):
    """
    Computes multi-head attention. Supports nested or padded tensors.

    Args:
        E_q (int): Size of embedding dim for query
        E_k (int): Size of embedding dim for key
        E_v (int): Size of embedding dim for value
        E_total (int): Total embedding dim of combined heads post input projection. Each head
            has dim E_total // nheads
        nheads (int): Number of heads
        dropout_p (float, optional): Dropout probability. Default: 0.0
    """
    def __init__(self, E_q: int, E_k: int, E_v: int, E_total: int,
                 nheads: int, dropout_p: float = 0.0):
        super().__init__()
        self.nheads = nheads
        self.dropout_p = dropout_p
        self.query_proj = nn.Linear(E_q, E_total)
        self.key_proj = nn.Linear(E_k, E_total)
        self.value_proj = nn.Linear(E_v, E_total)
        E_out = E_q
        self.out_proj = nn.Linear(E_total, E_out)
        assert E_total % nheads == 0, "Embedding dim is not divisible by nheads"
        self.E_head = E_total // nheads

    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
        """
        Forward pass; runs the following process:
            1. Apply input projection
            2. Split heads and prepare for SDPA
            3. Run SDPA
            4. Apply output projection

        Args:
            query (torch.Tensor): query of shape (N, L_t, E_q)
            key (torch.Tensor): key of shape (N, L_s, E_k)
            value (torch.Tensor): value of shape (N, L_s, E_v)

        Returns:
            attn_output (torch.Tensor): output of shape (N, L_t, E_q)
        """
        # Step 1. Apply input projection
        # TODO: demonstrate packed projection
        query = self.query_proj(query)
        key = self.key_proj(key)
        value = self.value_proj(value)

        # Step 2. Split heads and prepare for SDPA
        # reshape query, key, value to separate by head
        # (N, L_t, E_total) -> (N, L_t, nheads, E_head) -> (N, nheads, L_t, E_head)
        query = query.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2)
        # (N, L_s, E_total) -> (N, L_s, nheads, E_head) -> (N, nheads, L_s, E_head)
        key = key.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2)
        # (N, L_s, E_total) -> (N, L_s, nheads, E_head) -> (N, nheads, L_s, E_head)
        value = value.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2)

        # Step 3. Run SDPA
        # (N, nheads, L_t, E_head)
        attn_output = F.scaled_dot_product_attention(
            query, key, value, dropout_p=dropout_p, is_causal=True)
        # (N, nheads, L_t, E_head) -> (N, L_t, nheads, E_head) -> (N, L_t, E_total)
        attn_output = attn_output.transpose(1, 2).flatten(-2)

        # Step 4. Apply output projection
        # (N, L_t, E_total) -> (N, L_t, E_out)
        attn_output = self.out_proj(attn_output)

        return attn_output

按照Transformer论文设置超参数

N = 512
E_q, E_k, E_v, E_total = 512, 512, 512, 512
E_out = E_q
nheads = 8

除了dropout概率:设置为0以进行正确性检查

dropout_p = 0.0

让我们从Zipf定律生成一些逼真的假数据。

def zipf_sentence_lengths(alpha: float, batch_size: int) -> torch.Tensor:
    # generate fake corpus by unigram Zipf distribution
    # from wikitext-2 corpus, we get rank "." = 3, "!" = 386, "?" = 858
    sentence_lengths = np.empty(batch_size, dtype=int)
    for ibatch in range(batch_size):
        sentence_lengths[ibatch] = 1
        word = np.random.zipf(alpha)
        while word != 3 and word != 386 and word != 858:
            sentence_lengths[ibatch] += 1
            word = np.random.zipf(alpha)
    return torch.tensor(sentence_lengths)

创建嵌套张量批量输入

def gen_batch(N, E_q, E_k, E_v, device):
    # generate semi-realistic data using Zipf distribution for sentence lengths
    sentence_lengths = zipf_sentence_lengths(alpha=1.2, batch_size=N)

    # Note: the torch.jagged layout is a nested tensor layout that supports a single ragged
    # dimension and works with torch.compile. The batch items each have shape (B, S*, D)
    # where B = batch size, S* = ragged sequence length, and D = embedding dimension.
    query = torch.nested.nested_tensor([
        torch.randn(l.item(), E_q, device=device)
        for l in sentence_lengths
    ], layout=torch.jagged)

    key = torch.nested.nested_tensor([
        torch.randn(s.item(), E_k, device=device)
        for s in sentence_lengths
    ], layout=torch.jagged)

    value = torch.nested.nested_tensor([
        torch.randn(s.item(), E_v, device=device)
        for s in sentence_lengths
    ], layout=torch.jagged)

    return query, key, value, sentence_lengths

query, key, value, sentence_lengths = gen_batch(N, E_q, E_k, E_v, device)

生成用于比较的填充形式的查询、键、值

def jagged_to_padded(jt, padding_val):
    # TODO: do jagged -> padded directly when this is supported
    return torch.nested.to_padded_tensor(
        torch.nested.nested_tensor(list(jt.unbind())),
        padding_val)

padded_query, padded_key, padded_value = (
    jagged_to_padded(t, 0.0) for t in (query, key, value)
)

构建模型

mha = MultiHeadAttention(E_q, E_k, E_v, E_total, nheads, dropout_p).to(device=device)

检查正确性和性能

def benchmark(func, *args, **kwargs):
    torch.cuda.synchronize()
    begin = timeit.default_timer()
    output = func(*args, **kwargs)
    torch.cuda.synchronize()
    end = timeit.default_timer()
    return output, (end - begin)

output_nested, time_nested = benchmark(mha, query, key, value)
output_padded, time_padded = benchmark(mha, padded_query, padded_key, padded_value)

# padding-specific step: remove output projection bias from padded entries for fair comparison
for i, entry_length in enumerate(sentence_lengths):
    output_padded[i, entry_length:] = 0.0

print("=== without torch.compile ===")
print("nested and padded calculations differ by", (jagged_to_padded(output_nested, 0.0) - output_padded).abs().max().item())
print("nested tensor multi-head attention takes", time_nested, "seconds")
print("padded tensor multi-head attention takes", time_padded, "seconds")

# warm up compile first...
compiled_mha = torch.compile(mha)
compiled_mha(query, key, value)
# ...now benchmark
compiled_output_nested, compiled_time_nested = benchmark(
    compiled_mha, query, key, value)

# warm up compile first...
compiled_mha(padded_query, padded_key, padded_value)
# ...now benchmark
compiled_output_padded, compiled_time_padded = benchmark(
    compiled_mha, padded_query, padded_key, padded_value)

# padding-specific step: remove output projection bias from padded entries for fair comparison
for i, entry_length in enumerate(sentence_lengths):
    compiled_output_padded[i, entry_length:] = 0.0

print("=== with torch.compile ===")
print("nested and padded calculations differ by", (jagged_to_padded(compiled_output_nested, 0.0) - compiled_output_padded).abs().max().item())
print("nested tensor multi-head attention takes", compiled_time_nested, "seconds")
print("padded tensor multi-head attention takes", compiled_time_padded, "seconds")

请注意,如果没有torch.compile,Python子类嵌套张量的开销可能会使其比在填充张量上的等效计算更慢。然而,一旦启用torch.compile,在嵌套张量上的操作会带来多倍的加速。随着批次中填充比例的增加,避免在填充上的浪费计算变得更加有价值。

print(f"Nested speedup: {compiled_time_padded / compiled_time_nested:.3f}")

结论

在本教程中,我们学习了如何使用嵌套张量执行基本操作,以及如何以避免在填充上进行计算的方式实现变压器的多头注意力机制。更多信息,请查看torch.nested命名空间的文档。

脚本总运行时间: ( 0 分钟 0.000 秒)

Gallery generated by Sphinx-Gallery

优云智算