Torch 常用方法记录

summary: 记录一些常用torch方法争取做到可以默写下来~ tags: export category: 编程

因为感觉torch里面实现相同功能可以有很多种不同方法这里希望约定俗成一下

让自己习惯每次只用同一种方法不用再去反复查语法确认

1. 复制与分割 unsqueeze + repeat + unbind

#1. 复制与分割 unsqueeze + repeat + unbind
# input with dim(5,3) 分成3份 这里dim=0 会默认把第一个dim每个元素分别提出来
input = torch.randn(5,3)

input = input.unsqueeze(0).repeat(3,1,1)
a, b, c = torch.unbind(input, dim = 0)

print(*map(lambda x: x.shape, [a,b,c])

2. 区分einops.rearrange 和 torch.einsum

#2. 区分einops.rearrange 和 torch.einsum
#   rearrange 处理matrix内部结构 替代常见torch command(view, reshape, transpose, unsqueeze, ...)
#   注意rearrange 每个dim必须用空格隔开 因为允许 长字符串代表dim 而einsum中只能单字符
#   并且rearrange 允许括号 einsum中没有internally change dim的需求 因此不允许
#   感觉比细分command好用因为dim会更加容易理解 时刻可以知道是哪几个dim互相交换和当前dim

# (batch_size, channels, height, width)
x = torch.randn(16, 3, 32, 32)

x_reshaped = x.view(x.shape[0], -1)
x_reshaped = rearrange(x, 'b c h w -> b (c h w)')

x_transposed = x.permute(0, 1, 3, 2)
x_transposed = rearrange(x, 'b c h w -> b c w h')

x_unsqueezed = x.unsqueeze(1)  # dim=1位unsqueeze 加dim
x_squeezed = x_unsqueezed.squeeze()  #不标注dim 去除所有值为1的dim

x_unsqueezed = rearrange(x, 'b c h w -> b 1 c h w')
x_squeezed = rearrange(x_unsqueezed, 'b 1 c h w -> b c h w')

#使用场景-attention
def self_attend(x, num_head = 8):

    # 假设x dim为(b, l, d) -> (b, l, (n, -1))
    x = rearrange(x, 'b l (n d) -> b n l d', n = num_head)

    q,k,v = torch.unbind(x.unsqueeze(0).repeat(3,1,1), dim = 0)

 
    scores = torch.einsum('b n q c, b n k c -> b n q k', q, k) / ((x.shape[-1]) ** 0.5)
    weights = torch.softmax(scores, dim = -1)

    attend = torch.einsum('b n q k, b n k d -> b n q d', weights, v)

    return rearrange(attend, 'b n q d -> b q (n d)', n = num_head)

#使用场景-Covariance
def cov(a):

    mean = torch.mean(a, dim=0)
    a = a - mean

    return torch.einsum('ba, bc -> ac', a, a) / a.shape[0]

#使用场景-determinant
A = torch.randn(batch_size, n, n)

det = torch.einsum('bii->b', A)

# Polynomial Feature Expansion
X = torch.randn(batch_size, features)

# Compute all pairwise products (second-degree polynomial features)
poly_features = torch.einsum('bi,bj->bij', X, X).view(batch_size, -1)

# Equivalent to torch.matmul(X.unsqueeze(2), X.unsqueeze(1)).view(batch_size, -1)

1. 复制与分割 unsqueeze + repeat + unbind

2. 区分einops.rearrange 和 torch.einsum

3. F.pad