一、前言
在深挖ML4CO的代码过程中，遇到了torch.take_along_dim()这个方法，影响到我后续的代码阅读；加之在上网搜索资料的过程中，网络上对此函数的介绍文章少之又少，即使有，也是对torch官网文档中的解释进行英译汉或对函数的轻描淡写，看完解析后，对该函数的认识并没有因此而深刻。故在此谈一下自己的理解。
温馨提示：由于torch.take_along_dim函数与torch.gather函数的功能大同小异，因此在阅读这篇文章前，建议先阅读torch.gather()函数的解析PyTorch基础（16）-- torch.gather()。
二、方法解析
首先，我们需要看一下torch官方对于该函数的解释。白话解释为：我有一个tensor为input，沿着给定的input的维度（dim），给定取值的索引（indices），从input取出我想要的值。嗯……跟torch.gather方法不能说一模一样，只能说完全相同啊。
这个函数如何使用，直接通过案例来解释。
三、案例分析
3.1 案例1（不指定dim）
值得注意的是：不指定dim时，该方法会将input展开为一个一维tensor，然后根据indices进行取值。
import torch
t = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
index = torch.tensor([[2,0,2],[1,2,0],[1,1,1]])
t1 = torch.take_along_dim(t, index)
t1
输出为：
3.2 案例2（dim=0）
import torch
t = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
index = torch.tensor([[2,0,2],[1,2,0],[1,1,1]])
t1 = torch.take_along_dim(t, index, dim=0)
t1
3.3 案例3（dim=1）
import torch
t = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
index = torch.tensor([[2,0,2],[1,2,0],[1,1,1]])
t1 = torch.take_along_dim(t, index, dim=1)
t1
参考文献
- https://pytorch.org/docs/stable/generated/torch.take_along_dim.html#torch.take_along_dim
- https://blog.csdn.net/dongjinkun/article/details/132299818?spm=1001.2014.3001.5501
