基于torch.where和布尔索引的速度比较

技术  /  管理员 发布于 7年前   232

我就废话不多说了，直接上代码吧！

import torchimport timex = torch.Tensor([[1, 2, 3], [5, 5, 5], [7, 8, 9],[5,5,5],[1,2,3,],[1,2,4]])'''使用pytorch实现对于任意shape的torch.tensor,如果其中的element不等于5则为0，等于5则保留原数值实现该功能的两种方式，并比较两种实现方式的速度'''# x[x!=5]=1def t2(x):  x[x!=5]=0  return xdef t(x):  zeros=torch.zeros(x.shape)  # ones=torch.ones(x.shape)  x=torch.where(x!=5,zeros,x)  return xt2_start=time.time()t2=t2(x)t2_end=time.time()t_start=time.time()t=t(x)t_end=time.time()print(t2,t)print(torch.sum(t-t2))print('using x[x!=5]=0 time:',t2_end-t2_start)print('using torch.where time:',t_end-t_start)'''tensor([[0., 0., 0.],    [5., 5., 5.],    [0., 0., 0.],    [5., 5., 5.],    [0., 0., 0.],    [0., 0., 0.]]) tensor([[0., 0., 0.],    [5., 5., 5.],    [0., 0., 0.],    [5., 5., 5.],    [0., 0., 0.],    [0., 0., 0.]])tensor(0.)using x[x!=5]=0 time: 0.0010008811950683594using torch.where time: 0.0看来大神说的没错，果然是使用torch.where速度更快 a[a!=5]=0 这种写法，速度比 torch.where 慢了超级多'''

以上这篇基于torch.where和布尔索引的速度比较就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持。