torch.jit.trace能否正确记录ReLU操作？

1.答案

能。

2.验证过程

由于torch.jit.trace的开发文档中写出无法正确记录控制流的代码，但是nn.ReLU操作它明显是个控制流的代码，所以它能够被正确记录吗？
下面生成一个简单的由Linear+ReLU操作组成的网络验证下。

# 验证relu操作是否会被torch.jit.trace正确跟踪
import torch
from torch import nn


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(1, 1, bias=False)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.fc(x)
        x = self.relu(x)
        return x


net = Net()
net.eval()
net.fc.weight.data = torch.ones(1)

a = torch.ones(1, 1)
a_out = net(a)
print(a_out)  # 输出1
b = -torch.ones(1, 1)
b_out = net(b)  # 输出0
print(b_out)

script_module = torch.jit.trace(net, a)
print(script_module(b))  # 输出0，可以分析出torch.jit.trace可以正确跟踪relu操作

最终输出的是：

tensor([1.], grad_fn=<ReluBackward0>)
tensor([0.], grad_fn=<ReluBackward0>)
tensor([0.], grad_fn=<ReluBackward0>)

可以看出，即使我们是以正数a作为数据样例让torch.jit.trace跟踪模型，得到的script_module依然能够正确将负数b输出为0.可以看到relu操作起到了作用。