PyTorch中的take_along_dim
技术背景在此前的一篇博客中，我们介绍过take_along_axis这个算子的具体使用方法。这里针对于Pytorch的take_along_dim算子，再重新介绍一次。
Numpy版本使用
这里我们展示的案例是基于numpy-2.0.1版本实现的：
$ python3 -m pip show numpy
Name: numpy
Version: 2.0.1
Summary: Fundamental package for array computing in Python
Home-page: https://numpy.org
Author: Travis E. Oliphant et al.示例如下：
In : import numpy as np
In : a = np.arange(12).reshape((1,4,3))
In : a
Out:
array([[[ 0,1,2],
[ 3,4,5],
[ 6,7,8],
[ 9, 10, 11]]])
In : idx = np.array()
In : b = np.take_along_axis(a, idx, axis=1)
In : b
Out:
array([[,
]])
In : b = np.take_along_axis(a, idx, axis=2)
In : b
Out:
array([[[ 1,2],
[ 4,5],
[ 7,8],
]])在这个基础示例中，我们分别展示了同一个索引矩阵，在不同的维度上进行索引的结果。使用take_along_axis有一个默认的要求：原始数组和索引数组的维度数量需要保持一致。但是因为这里的索引矩阵是一维的，那么我们只要用slice的方法对索引矩阵进行扩维就好了。例如，我们需要在第二个维度进行提取，那么就可以用arr来进行扩维。
PyTorch版实现
这里我们使用的torch是2.5.1的稳定版：
$ python3 -m pip show torch
Name: torch
Version: 2.5.1
Summary: Tensors and Dynamic neural networks in Python with strong GPU acceleration
Home-page: https://pytorch.org/
Author: PyTorch Team
Author-email: packages@pytorch.org
License: BSD-3-Clause
Location: /miniconda3/envs/pytorch/lib/python3.9/site-packages
Requires: filelock, fsspec, jinja2, networkx, sympy, typing-extensions
Required-by: torchaudio, torchmetrics, torchvision相关的API接口文档如下：
其实实现起来跟numpy的操作是非常类似的：
In : import torch as tc
In : a = tc.arange(12).reshape((1,4,3))
In : idx = tc.tensor()
In : b = tc.take_along_dim(a, idx, dim=1)
In : b
Out:
tensor([[,
]])
In : b = tc.take_along_dim(a, idx, dim=2)
In : b
Out:
tensor([[[ 1,2],
[ 4,5],
[ 7,8],
]])可以说是基本一致。那么同样的，也是要做一个扩维的处理。唯一一个不同的地方就是，在torch中是take_along_dim而不是像numpy或者mindspore中的take_along_axis，在torch中用dim替代了axis，包括函数名称和传入的关键词参数。
总结概要
接前面一篇take_along_axis的文章，本文主要介绍在PyTorch框架下，功能基本一样的函数take_along_dim。二者除了命名和一些关键词参数不一致之外，用法是一样的。需要注意的是，两者都要求输入的数组和索引数组维度数量一致。在特定场景下，需要手动进行扩维。
版权声明
本文首发链接为：https://www.cnblogs.com/dechinphy/p/take_along_dim.html
作者ID：DechinPhy
更多原著文章请参考：https://www.cnblogs.com/dechinphy/
打赏专用链接：https://www.cnblogs.com/dechinphy/gallery/image/379634.html
腾讯云专栏同步：https://cloud.tencent.com/developer/column/91958
来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！
页:
[1]