昇腾社区首页
中文
注册

beta)torch_npu.contrib.BiLSTM

该接口计划废弃,可以参考小算子拼接方案进行替换。

接口原型

torch_npu.contrib.BiLSTM(input_size, hidden_size)

功能描述

将NPU兼容的双向LSTM操作应用于输入序列。

参数说明

  • input_size:对输入期望的特征数量。
  • hidden_size:hidden state中的特征数量。

支持的型号

  • Atlas 训练系列产品
  • Atlas A2 训练系列产品
  • Atlas A3 训练系列产品
  • Atlas 推理系列产品

调用示例

1
2
3
>>> r = torch_npu.contrib.BiLSTM(512, 256).npu()
>>> input_tensor = torch.randn(26, 2560, 512).npu()
>>> output = r(input_tensor)