如何迭代PyTorch张量

我有一个大小为（1000110）的张量数据，我想迭代张量的第一个索引，并计算以下结果

    data = torch.randn(size=(1000,110)).to(device)
    
    male_poor = torch.tensor(0).float().to(device)
    male_rich = torch.tensor(0).float().to(device)
    
    female_poor = torch.tensor(0).float().to(device)
    female_rich = torch.tensor(0).float().to(device)
    
    for i in data:
    
        if torch.argmax(i[64:66]) == 0 and torch.argmax(i[108:110]) == 0:
          female_poor += 1
        if torch.argmax(i[64:66]) == 0 and torch.argmax(i[108:110]) == 1:
          female_rich += 1
        if torch.argmax(i[64:66]) == 1 and torch.argmax(i[108:110]) == 0:
          male_poor += 1
        if torch.argmax(i[64:66]) == 1 and torch.argmax(i[108:110]) == 1:
          male_rich += 1


    disparity = ((female_rich/(female_rich + female_poor))) / ((male_rich/(male_rich + male_poor)))

有没有比for loop更快的方法来实现这一点？

pytorch（以及numpy）中的关键是矢量化，也就是说，如果可以通过对矩阵进行操作来删除循环，那么速度会快得多。与底层编译的C代码中的循环相比，python中的循环速度非常慢。在我的机器上，代码的执行时间约为0.091s，以下矢量化代码约为0.002s，因此大约快x50：

import torch
torch.manual_seed(0)
device = torch.device('cpu')

data = torch.randn(size=(1000, 110)).to(device)

import time
t = time.time()
#vectorize over first dimension
argmax64_0 = torch.argmax(data[:, 64:66], dim=1) == 0
argmax64_1 = torch.argmax(data[:, 64:66], dim=1) == 1
argmax108_0 = torch.argmax(data[:, 108:110], dim=1) == 0
argmax108_1 = torch.argmax(data[:, 108:110], dim=1) == 1
female_poor = (argmax64_0 & argmax108_0).sum()
female_rich = (argmax64_0 & argmax108_1).sum()
male_poor = (argmax64_1 & argmax108_0).sum()
male_rich = (argmax64_1 & argmax108_1).sum()

disparity = ((female_rich / (female_rich + female_poor))) / ((male_rich / (male_rich + male_poor)))

print(time.time()-t)
print(disparity)

pytorch（以及numpy）中的关键是矢量化，也就是说，如果您可以通过对矩阵进行操作来删除循环，那么速度会快得多。与底层编译的C代码中的循环相比，python中的循环速度非常慢。在我的机器上，代码的执行时间约为0.091s，以下矢量化代码约为0.002s，因此大约快x50：

import torch
torch.manual_seed(0)
device = torch.device('cpu')

data = torch.randn(size=(1000, 110)).to(device)

import time
t = time.time()
#vectorize over first dimension
argmax64_0 = torch.argmax(data[:, 64:66], dim=1) == 0
argmax64_1 = torch.argmax(data[:, 64:66], dim=1) == 1
argmax108_0 = torch.argmax(data[:, 108:110], dim=1) == 0
argmax108_1 = torch.argmax(data[:, 108:110], dim=1) == 1
female_poor = (argmax64_0 & argmax108_0).sum()
female_rich = (argmax64_0 & argmax108_1).sum()
male_poor = (argmax64_1 & argmax108_0).sum()
male_rich = (argmax64_1 & argmax108_1).sum()

disparity = ((female_rich / (female_rich + female_poor))) / ((male_rich / (male_rich + male_poor)))

print(time.time()-t)
print(disparity)

---

我在argmax文档中没有看到模糊的争论，谢谢你的帮助！我在argmax文档中没有看到模糊的争论，谢谢你的帮助！