一个选项是遮罩图层的渐变矩阵

layer = nn.Linear(d_in, d_out)
grad_mask = torch.eye(layer.weight.shape)

# stuff happens here, at some point loss.backward() is called and populates layer.weight.grad

layer.weight.grad = layer.weight.grad * grad_mask

字符串
也就是说，如果你只想要对角线，而你希望所有其他值都为零，为什么要使用线性层呢？你可以直接将输入乘以对角线值。

import torch
import torch.nn as nn

d_weight = 8
weight = torch.zeros(d_weight, d_weight)
diag_values = torch.randn(d_weight)
weight = torch.diagonal_scatter(weight, diag_values, 0)

layer = nn.Linear(d_weight, d_weight, bias=False)
layer.weight.data = weight.data

x = torch.randn(1,8)
y1 = layer(x)
y2 = x*diag_values

(y1 == y2).all()
>True

型
你可以把它作为一个模块

class MyLayer(nn.Module):
    def __init__(self, weight_tensor):
        super().__init__()
        self.weight_tensor = nn.Parameter(weight_tensor)
        
    def forward(self, x):
        return x * self.weight_tensor

型