PyTorch torch.bmm 函数
torch.bmm 是 PyTorch 中用于执行批量矩阵乘法的函数。
函数定义
torch.bmm(input, mat2, out)
使用示例
实例
import torch
# 批量矩阵乘法
batch_a = torch.randn(10, 3, 4)
batch_b = torch.randn(10, 4, 5)
result = torch.bmm(batch_a, batch_b)
print("批量结果形状:", result.shape)
# 批量矩阵乘法
batch_a = torch.randn(10, 3, 4)
batch_b = torch.randn(10, 4, 5)
result = torch.bmm(batch_a, batch_b)
print("批量结果形状:", result.shape)
输出结果为:
批量结果形状: torch.Size([10, 3, 5])

Pytorch torch 参考手册