代码场景:在目标检测中获取到了类别,置信度和bbox的tensor,但是我想只想要指定类别的数据集咋办,比如只要class=1的数据,于是便有了下面代码 import torch cls = torch.Tensor([1, 1, 2, 3, 1]) conf = torch.Tensor([0.1, 0.2, 0.3, 0.4, 0.5]) bbox = torch.Tensor([[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4], [5, 5, 5, 5]]) #print(bbox) print(conf [cls==1]) print(bbox [cls==1])
[pytorch][原创]获取torch tensor中指定条件的数据
最新推荐文章于 2023-12-25 07:38:54 发布