我的PyTorch转发功能可以执行额外的操作吗

Question

通常，forward函数将一堆层串在一起，并返回最后一个层的输出。在返回之前，我可以在最后一层之后做一些额外的处理吗？例如，通过.view的一些标量乘法和整形？

我知道autograd在某种程度上可以计算出梯度。所以我不知道我的额外处理是否会以某种方式搞砸。谢谢

Answer 1

pytorch通过张量的计算图跟踪梯度，而不是通过函数。只要你的张量具有requires_grad=True属性，而它们的grad不是None，你就可以(几乎(做任何你喜欢的事情，并且仍然能够反向投影
只要您使用pytorch的操作(例如，此处和此处列出的操作(，您就应该没事。

有关更多信息，请参阅此。

例如(取自torchvision的VGG实现(：

class VGG(nn.Module):
def __init__(self, features, num_classes=1000, init_weights=True):
super(VGG, self).__init__()
#  ...
def forward(self, x):
x = self.features(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)  # <-- what you were asking about
x = self.classifier(x)
return x

在torchvision的ResNet:实现中可以看到一个更复杂的例子

class Bottleneck(nn.Module):
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
base_width=64, dilation=1, norm_layer=None):
super(Bottleneck, self).__init__()
# ...
def forward(self, x):
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
if self.downsample is not None:    # <-- conditional execution!
identity = self.downsample(x)
out += identity  # <-- inplace operations
out = self.relu(out)
return out