位置: IT常识 - 正文
torch.where()用法(torch.save用法)
编辑：rootadmin推荐整理分享torch.where()用法(torch.save用法)，希望有所帮助，仅作参考，欢迎阅读内容。
文章相关热门搜索词:torch.sort用法,torch.save用法,torch.cat用法,torch.save用法,torch where,torch.where()函数,torch where,torch的用法,内容如对您有帮助，希望把文章链接给更多的朋友！
本文主要讲述torch.where()的两种用法，第一种是最常规的，也是官方文档所注明的；第二种就是配合bool型张量的计算
1、torch.where()常规用法我们先看官方文档的解释：
torch.where(condition, x, y) 根据条件，也就是condiction，返回从x或y中选择的元素的张量（这里会创建一个新的张量，新张量的元素就是从x或y中选的，形状要符合x和y的广播条件）。 Parameters解释如下： 1、condition (bool型张量) ：当condition为真，返回x的值，否则返回y的值 2、x (张量或标量)：当condition=True时选x的值 2、y (张量或标量)：当condition=False时选y的值
我看了好些博文，他们都说x和y的形状必须相同，完全胡扯嘛，官方文档写的明明白白的：The tensors condition, x, y must be broadcastable. 也就是说condition、x、y能进行广播就行，并不要求形状一样。下面看用法：
1.1 形状相同先演示形状相同的情况：
import torchx = torch.tensor([[1, 2, 3], [3, 4, 5], [5, 6, 7]])y = torch.tensor([[5, 6, 7], [7, 8, 9], [9, 10, 11]])z = torch.where(x > 5, x, y)print(f'x = {x}')print(f'=========================')print(f'y = {y}')print(f'=========================')print(f'x > 5 = {x > 5}')print(f'=========================')print(f'z = {z}')>print result:x = tensor([[1, 2, 3], [3, 4, 5], [5, 6, 7]])=========================y = tensor([[ 5, 6, 7], [ 7, 8, 9], [ 9, 10, 11]])=========================x > 5 = tensor([[False, False, False], [False, False, False], [False, True, True]])=========================z = tensor([[5, 6, 7], [7, 8, 9], [9, 6, 7]])
上面定义了x和y，两者的形状shape=(3, 3)相同，然后condition = x > 5也是就x中的每个元素值都要大于5，这里就能看到x中第0行和第1行都是False，只有第2行的1、2列是True，然后前面说了，为True时使用的是x中的值，为False时使用的是y中的值，那么新创建的z前两行和第2行0列使用的是y中的值，剩下两个使用x中的值，z的shape也是(3, 3)。
1.2 标量情况x = 3y = torch.tensor([[1, 5, 7]])z = torch.where(y > 2, y, x)print(f'y > 2 = {y > 2}')print(f'=========================')print(f'z = {z}')print(f'y > 2 = {y > 2}')print(f'=========================')print(f'z = {z}')>print result:y > 2 = tensor([[False, True, True]])=========================z = tensor([[3, 5, 7]])在这里，x是一个标量，condition = y > 2，你要是问我为什么不把condition设为condition = x > 2，很简单，x > 2不是bool Tensor。这里标量和张量是可以进行广播的！！ example:
a = torch.tensor([1, 5, 7])b = 3c = a + bd = torch.tensor([3, 3, 3])e = a + dprint(f'c = {c}')print(f'e = {e}')>print result:c = tensor([ 4, 8, 10])d = tensor([ 4, 8, 10])其实就是把b = 3拉成了[3, 3, 3]，也是就d那样。
1.3 形状不同其实标量那里也算是形状不同了，这里我再啰嗦一下吧，看例子：
x = torch.tensor([[1, 3, 5]])y = torch.tensor([[2], [4], [6]])z = torch.where(x > 2, x, y)print(f'x = {x}')print(f'=========================')print(f'y = {y}')print(f'=========================')print(f'x > 2 = {x > 2}')print(f'=========================')print(f'z = {z}')>print result:x = tensor([[1, 3, 5]])=========================y = tensor([[2], [4], [6]])=========================x > 2 = tensor([[False, True, True]])=========================z = tensor([[2, 3, 5], [4, 3, 5], [6, 3, 5]])上面x.shape=(1, 3) y.shape=(3, 1)，然后condition = x > 2的shape=(1, 3)，是可广播的，所以运算也能成功，在计算torch.where(x > 2, x, y)时，分别对x、y、condition进行广播，x.shape=(3, 3)，y.shape=(3, 3)，condition.shape=(3, 3) 所以y的值替换第0列，第1、2列为x的值。 更多的广播形式请读者朋友自行尝试
2、torch.where()特殊用法torch.where(a & b) a和b都是bool Tensor，返回的是一个元组，元组第一项是a、b中都为True的行的index的Tensor，第二项是a、b都为True列的index的Tensor
请看例子：
a = torch.tensor([[0, 1, 1], [1, 0, 0], [0, 0, 1]], dtype=torch.bool)b = torch.ones((3, 3), dtype=torch.bool)c = torch.where(a & b)print(f'a = {a}')print(f'=========================')print(f'b = {b}')print(f'=========================')print(f'c = {c}')>print result:a = tensor([[False, True, True], [ True, False, False], [False, False, True]])=========================b = tensor([[True, True, True], [True, True, True], [True, True, True]])=========================c = (tensor([0, 0, 1, 2]), tensor([1, 2, 0, 2]))c就是一个元组，第0项是a、b都为True的行标，第1项是a、b都为True的列标
总结以上就是torch.where()的两种用法，看起来比较麻烦，多练练也就是那样，特别一点的就是一个广播机制一个特殊用法，欢迎评论指正！ 请尊重原创，拒绝转载！！！
参考链接https://pytorch.org/docs/stable/generated/torch.where.html#torch.where https://pytorch.org/docs/stable/notes/broadcasting.html#broadcasting-semantics https://numpy.org/doc/stable/user/basics.broadcasting.html
上一篇：斯诺多尼亚国家公园多尔威泽兰城堡，英国威尔士康威 (© Sebastian Wasek/Sime/eStock Photo)(斯诺多尼亚山)
下一篇：雪花 (© TothGaborGyula/Getty Images Plus)
