Как Pytorch строит граф вычислений

Aug 25 2020

Вот пример кода pytorch с сайта:

class Net(nn.Module):

def __init__(self):
    super(Net, self).__init__()
    # 1 input image channel, 6 output channels, 3x3 square convolution
    # kernel
    self.conv1 = nn.Conv2d(1, 6, 3)
    self.conv2 = nn.Conv2d(6, 16, 3)
    # an affine operation: y = Wx + b
    self.fc1 = nn.Linear(16 * 6 * 6, 120)  # 6*6 from image dimension
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

def forward(self, x):
    # Max pooling over a (2, 2) window
    x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
    # If the size is a square you can only specify a single number
    x = F.max_pool2d(F.relu(self.conv2(x)), 2)
    x = x.view(-1, self.num_flat_features(x))
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x

В функции forward мы просто применяем серию преобразований к x, но никогда явно не определяем, какие объекты являются частью этого преобразования. Тем не менее, при вычислении градиента и обновлении весов Pytorch «волшебным образом» знает, какие веса обновлять и как следует рассчитывать градиент.

Как работает этот процесс? Проходит ли анализ кода или что-то еще, чего мне не хватает?

Ответы

1 AlexeyBirukov Aug 25 2020 at 23:08

Да, на прямом проходе есть неявный анализ. Изучите тензор результата, есть штука вроде grad_fn= <CatBackward>ссылки, позволяющая развернуть весь граф вычислений. И он построен в процессе реального прямого вычисления, независимо от того, как вы определили свой сетевой модуль, объектно-ориентированный с nn или «функциональным» способом.

Вы можете использовать этот график для анализа сети, как torchvizздесь:https://github.com/szagoruyko/pytorchviz/blob/master/torchviz/dot.py