tf.gather_nd 到火炬运营商答案

【问题标题】：tf.gather_nd to torch operatorstf.gather_nd 到火炬运营商
【发布时间】：2022-08-24 09:13:59
【问题描述】：

我看到了关于同一问题的主题，但我的要求在这里有所不同。我有这条线：

offsets = tf.gather_nd(offsets, kpt_inds, batch_dims=1)

而 offsets 是 (1,1,320,256,2) 大小的火炬张量，kpt_inds 是 (1,k,2) 的张量，k 是一个变量。

我想用一组能产生相同输出的火炬操作符来改变这个操作。该操作必须选择特定k 索引中的偏移量（在kpt_inds 中指定）。

我已经尝试过：

offsets = offsets[:, :, keypoints[:, :, 0], keypoints[:, :, 1], :]

它工作正常，但我有一个问题，我必须仅使用火炬运算符（没有 python 快捷方式）更改整个操作。原因是当我使用这个快捷方式时 tensorrt 的行为不正确。

标签： tensorflow torch tensorrt gather

【解决方案1】：

你可以试试.index_select 方法：

from einops import rearrange  # or use torch.unsqueeze instead
kpt_x = torch.ByteTensor(rearrange(keypoints[:, :, 0], '... -> 1 1 ...'))
kpt_y = torch.ByteTensor(rearrange(keypoints[:, :, 1], '... -> 1 1 1 ...'))
offsets = offsets.index_select(kpt_x)
offsets = offsets.index_select(kpt_y)

【讨论】：