我就廢話不多說了，直接上程式碼吧！

import torch
import time
x = torch.Tensor([[1,2,3],[5,5,5],[7,8,9],[1,3,],4]])
'''
使用pytorch實現對於任意shape的torch.tensor,如果其中的element不等於5則為0，等於5則保留原數值
實現該功能的兩種方式，並比較兩種實現方式的速度
'''

# x[x!=5]=1
def t2(x):
  x[x!=5]=0
  return x
def t(x):
  zeros=torch.zeros(x.shape)
  # ones=torch.ones(x.shape)
  x=torch.where(x!=5,zeros,x)
  return x

t2_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.]]) tensor([[0.,0.]])
tensor(0.)
using x[x!=5]=0 time: 0.0010008811950683594
using torch.where time: 0.0

看來大神說的沒錯，果然是使用torch.where速度更快
 a[a!=5]=0 這種寫法，速度比 torch.where 慢了超級多
'''

以上這篇基於torch.where和布林索引的速度比較就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。