2维张量在第一维度去重
时间: 2024-04-14 22:25:36 浏览: 152
如果你想在 PyTorch 中对一个二维张量在第一维度进行去重操作,可以使用 `torch.unique` 函数并指定 `dim=0` 参数来实现。
下面是一个示例代码,演示如何在第一维度对二维张量进行去重:
```python
import torch
# 创建一个二维张量
x = torch.tensor([[1, 2, 3], [4, 2, 1], [1, 2, 3]])
# 在第一维度进行去重
unique_values, indices = torch.unique(x, dim=0, return_inverse=True)
# 打印去重后的张量和索引
print("去重后的张量:")
print(unique_values)
print("原张量中每个元素在去重后张量中的索引:")
print(indices)
```
运行以上代码,会输出在第一维度去重后的张量和原张量中每个元素在去重后张量中的索引:
```
去重后的张量:
tensor([[1, 2, 3],
[4, 2, 1]])
原张量中每个元素在去重后张量中的索引:
tensor([0, 1, 0])
```
在这个示例中,我们使用 `torch.unique` 函数对二维张量进行去重操作,并指定 `dim=0` 参数来在第一维度进行去重。返回的 `unique_values` 张量包含了去重后的唯一值,而 `indices` 张量则表示了原张量中每个元素在去重后张量中的索引。
需要注意的是,返回的 `unique_values` 张量中的顺序是按照去重后的顺序排列的,而不是原始张量中的顺序。如果需要保留原始张量的顺序,可以使用 `torch.sort` 函数对索引进行排序操作。
阅读全文