（beta）torch_npu.npu_linear

接口原型

torch_npu.npu_linear(input, weight, bias=None) -> Tensor

功能描述

将矩阵“a”乘以矩阵“b”，生成“a*b”。

参数说明

• input (Tensor) - 2D矩阵张量。数据类型：float32、float16、int32、int8。格式：[ND, NHWC, FRACTAL_NZ]。
• weight (Tensor) - 2D矩阵张量。数据类型：float32、float16、int32、int8。格式：[ND, NHWC, FRACTAL_NZ]。
• bias (Tensor，可选，默认值为None) - 1D张量。数据类型：float32、float16、int32。格式：[ND, NHWC]。

调用示例

示例一：

>>> x=torch.rand(2,16).npu()
>>> w=torch.rand(4,16).npu()
>>> b=torch.rand(4).npu()
>>> output = torch_npu.npu_linear(x, w, b)
>>> output
tensor([[3.6335, 4.3713, 2.4440, 2.0081],
        [5.3273, 6.3089, 3.9601, 3.2410]], device='npu:0')

示例二：

from torch_npu.contrib.module.multihead_attention import NpuLinear
class NpuLinearModel(nn.Module):
    def __init__(self, input_size, output_size, bias_flag=True):
        super().__init__()
        self.layer = NpuLinear(input_size, output_size, bias_flag)
    def forward(self, x):
        return self.layer(x)

示例三：

class LinearApply(torch.autograd.Function):
    @staticmethod
    def forward(ctx, self, mat2, bias):
        self_view = self.view(-1, self.shape[-1])
        ctx.save_for_backward(self_view, mat2)
        ctx.linear_input_shape = self.shape
        out = torch_npu.npu_linear(self_view, mat2, bias)
        return out
    @staticmethod
    def backward(ctx, grad):
        self, mat2 = ctx.saved_tensors
        self_grad = torch_npu.npu_bmmV2(grad, mat2, []).view(ctx.linear_input_shape)
        mat2_grad = torch_npu.npu_bmmV2(grad.t(), self, [])
        bias_grad = grad.sum(0)
        return self_grad, mat2_grad, bias_grad
out = LinearApply.apply(total_input, weight, bias)