Interceptor

In this example, you gonna learn how to use Interceptor to capture intermediate values in the execution of a PyTorch module.

import torch # PyTorch

from pytorch_probing import Interceptor # Intercepts intermediate values.

We gonna start creating a example a module:

class ExampleModel(torch.nn.Module):

    def __init__(self, input_size, hidden_size, output_size, n_hidden=0):
        super().__init__()
        self.linear1 = torch.nn.Linear(input_size, hidden_size)
        self.relu = torch.nn.ReLU()

        if n_hidden > 0:
            layers = []
            for _ in range(n_hidden):
                layers.append(torch.nn.Linear(hidden_size, hidden_size))
                layers.append(torch.nn.ReLU())
            self.hidden_layers = torch.nn.Sequential(*layers)
        self._n_hidden = n_hidden

        self.linear2 = torch.nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.linear1(x)
        x = self.relu(x)
        if self._n_hidden > 0:
            x = self.hidden_layers(x)
        x = self.linear2(x)

        return x

In this case, we gonna use 2 hidden layers between the first and last layer:

input_size = 2
hidden_size = 3
output_size = 1
n_hidden = 2

model = ExampleModel(input_size, hidden_size, output_size, n_hidden)
model.eval()

ExampleModel(
  (linear1): Linear(in_features=2, out_features=3, bias=True)
  (relu): ReLU()
  (hidden_layers): Sequential(
    (0): Linear(in_features=3, out_features=3, bias=True)
    (1): ReLU()
    (2): Linear(in_features=3, out_features=3, bias=True)
    (3): ReLU()
  )
  (linear2): Linear(in_features=3, out_features=1, bias=True)
)

Than, we gonna pass our created module to a Interceptor, with the paths of the submodules we wanna get its outputs. Notices that we can use “.” to get inner submodules:

paths = ["linear1", "hidden_layers.2", "linear2"]

intercepted_model = Interceptor(model, paths, detach=False)
intercepted_model.eval()

Interceptor(
  (_module): ExampleModel(
    (linear1): InterceptorLayer(
      (_module): Linear(in_features=2, out_features=3, bias=True)
    )
    (relu): ReLU()
    (hidden_layers): Sequential(
      (0): Linear(in_features=3, out_features=3, bias=True)
      (1): ReLU()
      (2): InterceptorLayer(
        (_module): Linear(in_features=3, out_features=3, bias=True)
      )
      (3): ReLU()
    )
    (linear2): InterceptorLayer(
      (_module): Linear(in_features=3, out_features=1, bias=True)
    )
  )
)

The Interceptor modifies in-place the original module:

model

ExampleModel(
  (linear1): InterceptorLayer(
    (_module): Linear(in_features=2, out_features=3, bias=True)
  )
  (relu): ReLU()
  (hidden_layers): Sequential(
    (0): Linear(in_features=3, out_features=3, bias=True)
    (1): ReLU()
    (2): InterceptorLayer(
      (_module): Linear(in_features=3, out_features=3, bias=True)
    )
    (3): ReLU()
  )
  (linear2): InterceptorLayer(
    (_module): Linear(in_features=3, out_features=1, bias=True)
  )
)

We pass a example input throught the model:

inputs = torch.randn([10, 2])

with torch.no_grad():
    outputs = intercepted_model(inputs)

outputs

tensor([[-0.0029],
        [-0.0029],
        [-0.0029],
        [-0.0029],
        [-0.0029],
        [-0.0029],
        [-0.0029],
        [-0.0029],
        [-0.0029],
        [-0.0029]])

And the interceptor captures the required outputs and stores then in the “outputs” attribute:

intercepted_model.outputs

{'linear1': tensor([[ 0.1654,  0.0663, -0.7493],
         [ 0.3264, -0.9112, -0.8976],
         [ 0.1156,  0.0056, -0.8358],
         [ 0.7862, -1.3303, -0.4560],
         [ 0.3589, -1.0462, -0.9049],
         [ 0.0942, -0.3027, -0.9759],
         [ 0.7300, -1.1132, -0.4494],
         [ 0.5927, -0.4919, -0.4003],
         [ 0.3150, -0.3468, -0.7065],
         [-0.5145,  0.6613, -1.4115]]),
 'hidden_layers.2': tensor([[-0.5792, -0.1493, -0.1676],
         [-0.5853, -0.1682, -0.1807],
         [-0.5907, -0.1497, -0.1603],
         [-0.6053, -0.2082, -0.2035],
         [-0.5843, -0.1710, -0.1838],
         [-0.5923, -0.1483, -0.1580],
         [-0.6003, -0.2033, -0.2024],
         [-0.5879, -0.1913, -0.1997],
         [-0.5856, -0.1673, -0.1795],
         [-0.4808, -0.0937, -0.1802]]),
 'linear2': tensor([[-0.0029],
         [-0.0029],
         [-0.0029],
         [-0.0029],
         [-0.0029],
         [-0.0029],
         [-0.0029],
         [-0.0029],
         [-0.0029],
         [-0.0029]])}

We can clean this outputs with the interceptor_clear method:

intercepted_model.interceptor_clear()
intercepted_model.outputs

{'linear1': None, 'hidden_layers.2': None, 'linear2': None}

And return the model to its original state with the reduce method:

intercepted_model.reduce()

model

ExampleModel(
  (linear1): Linear(in_features=2, out_features=3, bias=True)
  (relu): ReLU()
  (hidden_layers): Sequential(
    (0): Linear(in_features=3, out_features=3, bias=True)
    (1): ReLU()
    (2): Linear(in_features=3, out_features=3, bias=True)
    (3): ReLU()
  )
  (linear2): Linear(in_features=3, out_features=1, bias=True)
)