删除pytorch中张量字典列表中所有张量的某些元素

Question

我是 pytorch 的新手，正在尝试通过玩简单的代码来学习。我有一个长度为零的列表，由字典填充，其中字典中的值是张量。这是这个列表的样子：

A = [{'boxes': tensor([[ 142.1232,  142.9373, 1106.0452,  971.3792],
    [ 259.1277,  618.4834, 1100.1293, 1028.8989],
    [ 232.1346,  692.5888,  763.3408, 1028.6766],
    [ 206.8070,  312.2080, 1137.1434, 1013.4373],
    [ 495.9471,  675.7287,  978.5932, 1012.7568]], grad_fn=<StackBackward>), 'labels': tensor([16,  1,  1,  1,  1]), 'scores': tensor([0.9988, 0.9489, 0.5228, 0.3500, 0.0639], grad_fn=<IndexBackward>), 'masks': tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]]], grad_fn=<UnsqueezeBackward0>)}]

这个列表是 Mask-RCNN 的输出，我想删除嵌套字典中所有张量的某些元素。在这种情况下，我只想保留与 class 标签“1”相关的信息。与每个 class（'boxes'、'labels'、'scores' 和 'masks'）相关的信息都位于每个中的相同位置（索引）张量。所以，我找到了所有“1"s in the tensor with key "标签”的索引：

idxOfClass = [i for i, x in enumerate(list(pred[0]['labels'])) if x == 1]

这给了我：[1, 2, 3, 4]。然后，我想在嵌套字典的所有张量中保留位于 idxOfClass 索引处的所有值。如果我这样做：

Anew = [{pred[0]['boxes'][idxOfClass],pred[0]['labels'][idxOfClass],pred[0]['masks'][idxOfClass],pred[0]['scores'][idxOfClass]}]

我得到：

[{tensor([[ 259.1277,  618.4834, 1100.1293, 1028.8989],
    [ 232.1346,  692.5888,  763.3408, 1028.6766],
    [ 206.8070,  312.2080, 1137.1434, 1013.4373],
    [ 495.9471,  675.7287,  978.5932, 1012.7568]], grad_fn= 
    <IndexBackward>), tensor([0.9489, 0.5228, 0.3500, 0.0639], 
    grad_fn= 
    <IndexBackward>), tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]],


    [[[0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      ...,
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.],
      [0., 0., 0.,  ..., 0., 0., 0.]]]], grad_fn=<IndexBackward>), 
    tensor([1, 1, 1, 1])}]

但是，这不是以张量作为字典值的字典列表。这是一个张量列表，没有嵌套字典的键值结构。我的问题是，"is there any way to keep the original structure of the list when I remove certain elements of all the tensors using indices of specific elements?".

Answer 1

您可以在构建新的预测结果时只添加键名。

Anew = [{'boxes': pred[0]['boxes'][idxOfClass],'labels': pred[0]['labels'][idxOfClass],'masks': pred[0]['masks'][idxOfClass],'scores': pred[0]['scores'][idxOfClass]}]

删除pytorch中张量字典列表中所有张量的某些元素

Remove certain elements of all the tensors in a list of dictionary of tensors in pytorch

python

dictionary

list

pytorch

tensor