删除 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 的输出，我想删除嵌套字典中所有张量的某些元素。在这种情况下，我只想保留与类标签"1"相关的信息。与每个类相关的信息("框"、"标签"、"分数"和"掩码"(都位于每个张量中的同一位置(索引(。因此，我找到了张量中所有"1"的索引，并带有键"标签"：

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])}]

但是，这不是一个以张量作为字典值的字典列表。这是一个张量列表，没有嵌套字典的键值结构。我的问题是，"当我使用特定元素的索引删除所有张量的某些元素时，有没有办法保持列表的原始结构？

Answer 1

您可以在构建新的 pred ress 时添加键名。

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