自定义层实现

要设计一个接受输入并计算张量降维的层，返回值为

$$y_k = \sum_{i,j} W_{ijk} x_i x_j$$

，我们需要实现一个自定义的 PyTorch 模块。这个层将接收一个输入张量 ( X )（通常是一维向量），并通过权重张量 ( W ) 计算输出。

实现步骤

1. 定义自定义层：该层将包含权重张量 ( W )。
2. 计算输出：在 forward 方法中，计算每个 ( y_k ) 的值。
3. 支持批处理：确保层可以处理多个输入样本。

以下是具体的实现代码：

自定义层实现

import torch
import torch.nn as nn

class TensorReductionLayer(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        # 初始化权重 W，大小为 (output_dim, input_dim, input_dim)
        self.W = nn.Parameter(torch.randn(output_dim, input_dim, input_dim) * 0.01)

    def forward(self, X):
        # X 的形状为 (batch_size, input_dim)
        # W 的形状为 (output_dim, input_dim, input_dim)
        
        # 计算 x_i * x_j 的外积
        outer_product = torch.matmul(X.unsqueeze(2), X.unsqueeze(1))  # 形状为 (batch_size, input_dim, input_dim)
        
        # 对每个输出维度计算加权求和
        # y_k 的形状为 (batch_size, output_dim)
        y_k = torch.einsum('bxy,zyx->bz', outer_product, self.W)  # batch_size, output_dim
        
        return y_k

# 示例使用
input_dim = 5   # 输入特征的维度
output_dim = 3  # 输出特征的维度
layer = TensorReductionLayer(input_dim, output_dim)

# 创建一个随机输入张量，大小为 (2, 5)
input_tensor = torch.rand(2, input_dim)

# 通过层进行前向传播
output_tensor = layer(input_tensor)

print("输出张量:\n", output_tensor)
print("输出张量的形状:", output_tensor.shape)  # 应该是 (2, output_dim)

代码解析

1. 权重 ( W )：

   • 在 __init__ 方法中，使用 nn.Parameter 来定义权重张量 ( W )，其形状为 ( (output_dim, input_dim, input_dim) )。这使得每个输出维度都可以根据输入特征的外积进行加权。

2. 前向传播：

   • 在 forward 方法中，通过 X.unsqueeze(2) 和 X.unsqueeze(1) 计算输入的外积，得到形状为 ( (batch_size, input_dim, input_dim) ) 的张量。
   • 使用 torch.einsum 进行加权求和，计算每个 ( y_k )，其形状为 ( (batch_size, output_dim) )。

预期结果

运行示例代码后，你应该能够看到输出张量和其形状的信息：

输出张量:
 tensor([...], grad_fn=<SumBackward0>)
输出张量的形状: torch.Size([2, output_dim])

这里的输出张量包含了每个输入样本的降维结果，形状为 ( (2, 3) )（假设 output_dim 设置为3）。

要理解 outer_product = torch.matmul(X.unsqueeze(2), X.unsqueeze(1))，我们可以通过一个数学示例来说明如何使用这个操作计算外积。外积是将两个向量结合起来生成一个矩阵。

设定示例

假设我们有一个输入张量 ( X )，其形状为 ( (2, 3) )，表示有 2 个样本，每个样本有 3 个特征。我们可以设定如下：

$$X = \begin{bmatrix} 1 & 2 & 3 \\ 4 & 5 & 6 \end{bmatrix}$$

使用 unsqueeze 操作

首先，使用 X.unsqueeze(2) 和 X.unsqueeze(1)：

• X.unsqueeze(2)：在第三维添加一个新维度。结果将变为形状为 ( (2, 3, 1) ) 的张量。

$$X.unsqueeze(2) = \begin{bmatrix} \begin{bmatrix} 1 \\ 2 \\ 3 \end{bmatrix} \end{bmatrix} \text{ 和 } \begin{bmatrix} \begin{bmatrix} 4 \\ 5 \\ 6 \end{bmatrix} \end{bmatrix}$$

• X.unsqueeze(1)：在第二维添加一个新维度。结果将变为形状为 ( (2, 1, 3) ) 的张量。

$$X.unsqueeze(1) = \begin{bmatrix} \begin{bmatrix} 1 & 2 & 3 \end{bmatrix} \end{bmatrix} \text{ 和 } \begin{bmatrix} \begin{bmatrix} 4 & 5 & 6 \end{bmatrix} \end{bmatrix}$$

计算外积

使用 torch.matmul 进行矩阵乘法：

import torch

X = torch.tensor([[1, 2, 3], [4, 5, 6]])
outer_product = torch.matmul(X.unsqueeze(2), X.unsqueeze(1))

• 外积的结果：此操作会将 ( X.unsqueeze(2) ) 和 ( X.unsqueeze(1) ) 进行矩阵乘法，从而生成一个形状为 ( (2, 3, 3) ) 的张量。

$$\text{outer_product} = \begin{bmatrix} \begin{bmatrix} 1 \times 1 & 1 \times 2 & 1 \times 3 \\ 2 \times 1 & 2 \times 2 & 2 \times 3 \\ 3 \times 1 & 3 \times 2 & 3 \times 3 \end{bmatrix} & \begin{bmatrix} 4 \times 1 & 4 \times 2 & 4 \times 3 \\ 5 \times 1 & 5 \times 2 & 5 \times 3 \\ 6 \times 1 & 6 \times 2 & 6 \times 3 \end{bmatrix} \end{bmatrix}$$

计算结果为：

$$\text{outer_product} = \begin{bmatrix} \begin{bmatrix} 1 & 2 & 3 \\ 2 & 4 & 6 \\ 3 & 6 & 9 \end{bmatrix} & \begin{bmatrix} 4 & 8 & 12 \\ 5 & 10 & 15 \\ 6 & 12 & 18 \end{bmatrix} \end{bmatrix}$$

结果解释

• 外积的每个切片（第一个维度的每个样本）将生成一个 ( 3 \times 3 ) 的矩阵，其中的每个元素 ( (i,j) ) 是 ( x_i ) 和 ( x_j ) 的乘积。
• 例如，对于第一个样本：

$$\begin{bmatrix} 1 & 2 & 3 \\ 2 & 4 & 6 \\ 3 & 6 & 9 \end{bmatrix}$$

表示的就是第一个样本特征的外积。

对于第二个样本：

$$\begin{bmatrix} 4 & 8 & 12 \\ 5 & 10 & 15 \\ 6 & 12 & 18 \end{bmatrix}$$

表示的就是第二个样本特征的外积。

总结

通过 unsqueeze 操作，我们将原始向量转化为适合进行外积计算的格式。torch.matmul 将两个张量相乘，从而得到外积结果，生成的张量的形状是 ( (2, 3, 3) )，其中包含了每个样本的外积矩阵。如果你有任何进一步的问题或需要更多例子，请随时告诉我！