Особенности определения моделей при переходе от Tensorflow к PyTorch

Question

Я только недавно переключился на PyTorch после того, как разочаровался в отладке tf, и понял, что это практически эквивалентно кодированию в numpy. Мой вопрос заключается в том, какие допустимые аспекты Python мы можем использовать в модели PyTorch (чтобы полностью поместить в GPU), например. if-else должно быть реализовано в тензорном потоке следующим образом

a = tf.Variable([1,2,3,4,5], dtype=tf.float32)
b = tf.Variable([6,7,8,9,10], dtype=tf.float32)
p = tf.placeholder(dtype=tf.float32)
ps = tf.placeholder(dtype=tf.bool)

li = [None]*5
li_switch = [True, False, False, True, True]

for i in range(5):
    li[i] = tf.Variable(tf.random.normal([5]))

sess = tf.Session()
sess.run(tf.global_variables_initializer())

def func_0():
    return tf.add(a, p)
def func_1():
    return tf.subtract(b, p)

with tf.device('GPU:0'):
    my_op = tf.cond(ps, func_1, func_0)

for i in range(5):
    print(sess.run(my_op, feed_dict={p:li[i], ps:li_switch[i]}))

Как изменится структура в pytorch для приведенного выше кода? Как разместить переменные и описанные выше операции на GPU и распараллелить входные данные списка для нашего графика в pytorch?

Answer 1

Чтобы инициализировать ваши тензоры a и b в PyTorch, вы делаете следующее:

a = torch.tensor([1,2,3,4,5], dtype=torch.float32)
b = torch.tensor([6,7,8,9,10], dtype=torch.float32)

Но , поскольку вам нужно, чтобы они были полностью на GPU, вы должны использовать волшебную функцию .cuda(). Итак, это будет:

a = torch.tensor([1,2,3,4,5], dtype=torch.float32).cuda()
b = torch.tensor([6,7,8,9,10], dtype=torch.float32).cuda()

Что перемещает тензор в графический процессор

---

Другой способ инициализации:

a = torch.FloatTensor([1,2,3,4,5]).cuda()
b = torch.FloatTensor([6,7,8,9,10]).cuda() 

---

Если нам нужно сгенерировать случайное нормальное распределение, мы используем torch.randn (есть также torch.rand, который делает равномерное случайное распределение).

li = torch.randn(5, 5)

(поймать ошибку, она должна быть инициализирована на cuda, вы не можете выполнять операции над тензорами, расположенными на отдельных процессорах, т. Е. CPU и GPU)

li = torch.randn(5, 5).cuda()

Нет никакой разницы для инициализации li_switch.

Один из возможных способов обработки func_0 и func_1 - объявить их как

def func_0(li_value):
    return torch.add(a, li_value)
def func_1(li_value):
    return torch.sub(b, li_value)

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

for i, pred in enumerate(li_switch):
    if pred:
        func_0(li[i])
    else:
        func_1(li[i])

---

Однако я предлагаю векторизовать ваши операции и сделать что-то вроде:

li_switch = torch.tensor([True, False, False, True, True])
torch.add(a, li[li_switch]).sum(dim=0)
torch.sub(b, li[~li_switch]).sum(dim=0)

Это гораздо более оптимизировано.

Answer 2

В pytorch код может быть написан так же, как и обычный код Python.

CPU

import torch
a = torch.FloatTensor([1,2,3,4,5])
b = torch.FloatTensor([6,7,8,9,10])
cond = torch.randn(5)

for ci in cond:
    if ci > 0:
        print(torch.add(a, 1))
    else:
        print(torch.sub(b, 1))

GPU

Переместите тензоры в GPU так:

a = torch.FloatTensor([1,2,3,4,5]).to('cuda')
b = torch.FloatTensor([6,7,8,9,10]).to('cuda')
cond = torch.randn(5).to('cuda')

import torch.nn as nn

class Cond(nn.Module):
    def __init__(self):
        super(Cond, self).__init__()

    def forward(self, cond, a, b):
        result =  torch.empty(cond.shape[0], a.shape[0]).cuda()
        for i, ci in enumerate(cond):
            if ci > 0:
                result[i] = torch.add(a, 1)
            else:
                result[i] = torch.sub(b, 1)

        return result

cond_model = Cond().to('cuda')
output = cond_model(cond, a, b)

https://pytorch.org/tutorials/beginner/blitz/tensor_tutorial.html#cuda-tensors