如何定义仅子对角块可训练的PyTorch方形张量？

你可以通过在forward方法中动态构建全零方阵并嵌入可训练子块来实现需求，无需在__init__中保存固定的全零张量（避免设备不匹配问题）。具体修改如下：

完整代码实现 #

import torch
import torch.nn as nn

class My_Network(nn.Module):
    def __init__(self, vertical_dim, horizontal_dim):
        super().__init__()
        self.vertical_dim = vertical_dim
        self.horizontal_dim = horizontal_dim
        self.total_dim = vertical_dim + horizontal_dim
        # 注册可训练的子对角块参数
        self.subdiagonal_block = nn.Parameter(torch.rand(vertical_dim, horizontal_dim))

    def forward(self):
        # 创建与可训练参数同设备、同数据类型的全零方阵
        full_matrix = torch.zeros(
            self.total_dim, self.total_dim,
            device=self.subdiagonal_block.device,
            dtype=self.subdiagonal_block.dtype
        )
        # 将子对角块嵌入矩阵的左下区域：
        # 行范围：从horizontal_dim到末尾（共vertical_dim行）
        # 列范围：从0到horizontal_dim（共horizontal_dim列）
        full_matrix[self.horizontal_dim:, :self.horizontal_dim] = self.subdiagonal_block
        return full_matrix

关键说明 #

• 动态构建全零矩阵：不在__init__中保存固定全零张量，而是在forward里实时生成，确保和可训练参数处于同一设备（CPU/GPU），避免后续计算时出现设备不匹配错误。
• 仅子块可训练：只有subdiagonal_block被nn.Parameter注册为模型参数，全零矩阵是临时生成的普通张量，不会被PyTorch的优化器更新。
• 切片定位：通过self.horizontal_dim:, :self.horizontal_dim精准定位左下区域，刚好匹配子块的vertical_dim × horizontal_dim尺寸。

使用示例 #

# 初始化模型：子块为3行2列，总方阵维度5×5
model = My_Network(vertical_dim=3, horizontal_dim=2)
full_tensor = model()
print(full_tensor.shape)  # 输出 torch.Size([5, 5])

# 验证可训练参数
print("可训练参数列表：")
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}")  # 仅输出 subdiagonal_block: torch.Size([3, 2])

# 用于矩阵乘法等计算
some_input = torch.rand(5, 10)
output = full_tensor @ some_input

内容的提问来源于stack exchange，提问作者Noumeno