问题
给出一个列表i_labels = [1,0,1,0],当取label ==1时,取images_true中的对应位置元素,当label == 0时,取images_false中对应位置的元素。如:
![](https://i-blog.csdnimg.cn/blog_migrate/11c3124f3bda0593bc4594551edb6a85.png)
解决方法1:torch.stack
images = torch.stack([images_true[idx] if label == 1 else images_false[idx] for idx, label in enumerate(i_labels)])
![](https://i-blog.csdnimg.cn/blog_migrate/221597d8194e3e04b83ca3dad123aac6.png)
解决方法2:对位相乘后相加
s = i_labels[:,None,None,None]
images = images * s + images_false * (1 - s)
images
![](https://i-blog.csdnimg.cn/blog_migrate/d5ec97a5b268e306306e67a050c2cb15.png)
比较
解决方法1与解决方法2的运行时间,均运行10W次,统计运行时间
![](https://i-blog.csdnimg.cn/blog_migrate/e0e8e65ec7e02d718ecd9169575ef699.png)
解决方法1 耗时7.06秒,解决方法2 耗时4.14秒。
总结
不同语句的运行时间是存在很大差距的,特别是用for语句。在我的30W数据集上跑,在一个epoch中使用解决方法2比解决方法1 ,可以节省2个小时。一条语句,省两个小时。
如您有更快的方式,欢迎讨论,让速度更快吧。