Подстановка по маске в PyTorch

Готовую маску можно использовать, чтобы собрать новый тензор: в ячейке с True берётся одно значение, с False - другое. Для этого служит функция where с тремя аргументами: маска, значение при истине, значение при лжи.

Создадим ряд, построим маску «больше 3» и подставим нули в остальные ячейки:

import torch

values = torch.tensor([1, 5, 3, 8])
mask = values > 3
result = torch.where(mask, values, torch.tensor(0))
print(result) # выведет tensor([0, 5, 0, 8])

Вместо нуля можно подставить любое число того же типа, что и исходный тензор:

import torch

temps = torch.tensor([2.0, 7.0, 4.0])
result = torch.where(temps > 5, temps, torch.tensor(-1.0))
print(result) # выведет tensor([-1.,  7., -1.])

Если при истине и при лжи передать разные тензоры одной формы, в каждой ячейке выберется элемент из нужного источника:

import torch

base = torch.tensor([10, 20, 30])
extra = torch.tensor([1, 2, 3])
mask = base > 15
result = torch.where(mask, base, extra)
print(result) # выведет tensor([ 1, 20, 30])

Создайте ряд из четырёх целых чисел и маску, где ячейка больше 5. Соберите новый ряд: там, где маска истинна, оставьте исходное число, иначе запишите -1, и выведите его.

Задайте ряд дробных температур и маску «выше 20». Соберите ряд, где в прохладных ячейках стоит 0.0, а в тёплых - исходное значение, и выведите результат.

Создайте два ряда одинаковой длины и маску равенства по позициям. Соберите ряд, где при совпадении берётся элемент первого ряда, иначе - второго, и выведите его.