溫馨提示×
您好，登錄后才能下訂單哦！
點擊 登錄注冊 即表示同意《億速云用戶服務條款》
您好，登錄后才能下訂單哦！
我就廢話不多說了，直接上代碼吧！
import torch
import time
x = 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]=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., 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.0010008811950683594
using torch.where time: 0.0
看來大神說的沒錯，果然是使用torch.where速度更快
a[a!=5]=0 這種寫法，速度比 torch.where 慢了超級多
'''
以上這篇基于torch.where和布爾索引的速度比較就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支持億速云。
免責聲明：本站發布的內容（圖片、視頻和文字）以原創、轉載和分享為主，文章觀點不代表本網站立場，如果涉及侵權請聯系站長郵箱：is@yisu.com進行舉報，并提供相關證據，一經查實，將立刻刪除涉嫌侵權內容。
