Skip to content

PyTorch Model Wrapper

interpretation.models.PyTorchModel

Bases: Model

This is a wrapper for PyTorch models

Source code in interpretation/models/pytorch_model.py 
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
class PyTorchModel(Model):
    """This is a wrapper for PyTorch models"""

    def __init__(
        self, 
        input_model: torch.nn.Module, 
        device: torch.device | str = None
    ) -> None:
        """Wraps a PyTorch model to the superclass"""

        if not isinstance(input_model, torch.nn.Module):
            raise TypeError("input_model must be an instance of torch.nn.Module")

        super().__init__(input_model)

        self.device = device or torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        self.model.to(self.device)
        self.model.eval()

        self.layer_dict = {
            name: module
            for name, module in self.model.named_modules()
            if len(list(module.children())) == 0
        }

    def __call__(self, X:npt.NDArray | torch.Tensor) -> npt.NDArray[Any]:
        "Model inference"
        dtype = next(self.model.parameters()).dtype

        with torch.no_grad():
            if isinstance(X, torch.Tensor):
                X_tensor = X.to(self.device, dtype=dtype)
            else:
                X_tensor = torch.tensor(X, device=self.device, dtype=dtype)

            output = self.model(X_tensor)

        return output.detach().cpu().numpy()

    def activation_at_layer(
        self, 
        X: npt.NDArray | torch.Tensor, 
        layer_identifier: int | str | list[int | str]
    ) -> list[npt.NDArray]:
        """
        Forward passes input X up to the specified layer.
        layer_identifier can be an integer index or layer name string.
        """ 
        if not isinstance(layer_identifier, list):
            layer_identifier = [layer_identifier]

        target_modules = {}
        all_modules = [m for m in self.model.modules()][1:]

        for out_idx, identifier in enumerate(layer_identifier):
            if isinstance(identifier, int):
                target_module = all_modules[identifier]
            elif isinstance(identifier, str):
                target_module = dict(self.model.named_modules()).get(identifier)

            if target_module is None:
                raise ValueError(f"Layers '{layer_identifier}' not found in PyTorch model.")

            target_modules[target_module] = out_idx

        activations = [None] * len(layer_identifier)
        handles = []

        def hook_generator(out_idx):
            def hook(module, input, output):
                activations[out_idx] = output.detach().cpu().numpy()
            return hook

        for module, out_idx in target_modules.items():
            handles.append(module.register_forward_hook(hook_generator(out_idx))) #runs the hook code when forward passed to the layer

        try:
            self.__call__(X)
        finally:
            for h in handles:
                h.remove()

        return activations

    def compute_gradients(
        self,
        X: npt.NDArray | torch.Tensor,
        objective: int | str | Tuple[int, int],
        objective_fn: Callable[[torch.Tensor], torch.tensor],
        wrt: int | str | Tuple[int, int] = None
    ) -> npt.NDArray:
        r"""Computes the gradient \(\frac{\partial \text{Objective}}{\partial X}\).

        `objective` is Objective, `wrt` is X.

        The gradient is calculated after applying the objective function to the objective layer w.r.t the wrt layer.

        Tuple forms of `objective` or `wrt` specifies specific neurons, e.g. (0, 1) means the second neuron of the first layer.

        Note: When no `wrt` is provided, it will default to `'input'`.
        """
        dtype = next(self.model.parameters()).dtype

        if isinstance(X, torch.Tensor):
            X_tensor = X.to(self.device, dtype=dtype)
        else:
            X_tensor = torch.tensor(X, device=self.device, dtype=dtype)        

        activations = {}
        hooks = []

        def hook_generator(name, neuron_idx=None):
            def hook(module, input, output):
                if neuron_idx is not None:
                    activations[name] = output.squeeze()[neuron_idx]
                else:
                    activations[name] = output
            return hook

        obj_name = self._layer_int_to_name(objective)
        wrt_name = self._layer_int_to_name(wrt) if wrt is not None else "input"

        if isinstance(objective, (int, str)):
            hooks.append(self.layer_dict[obj_name].register_forward_hook(hook_generator("obj")))
        else:
            hooks.append(self.layer_dict[obj_name].register_forward_hook(hook_generator("obj", objective[1])))

        if wrt_name != "input":
            hooks.append(self.layer_dict[wrt_name].register_forward_hook(hook_generator("wrt")))

        try:
            if wrt_name == "input":
                target_wrt = X_tensor
                target_wrt.requires_grad_()

                _ = self.model(target_wrt)
            else:
                def retain_grad_hook(module, input, output):
                    output.retain_grad()
                    activations["wrt"] = output

                hooks.append(self.layer_dict[wrt_name].register_forward_hook(retain_grad_hook))

                _ = self.model(X_tensor)
                target_wrt = activations["wrt"]

            obj_activation = activations["obj"]
            loss = objective_fn(obj_activation)

            if target_wrt.grad is not None:
                target_wrt.grad.zero_()
            loss.backward()

            grad = target_wrt.grad.detach().cpu().numpy()

            if isinstance(wrt, Tuple):
                return grad.squeeze()[wrt[1]]

            return grad

        finally:
            for h in hooks:
                h.remove()

    def _layer_int_to_name(self, layer_id: int | str | Tuple[int, int]) -> str:
        if isinstance(layer_id, int):
            return list(self.layer_dict.keys())[layer_id]
        elif isinstance(layer_id, Tuple):
            return list(self.layer_dict.keys())[layer_id[0]]
        return layer_id

__call__(X)

Model inference

Source code in interpretation/models/pytorch_model.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
def __call__(self, X:npt.NDArray | torch.Tensor) -> npt.NDArray[Any]:
    "Model inference"
    dtype = next(self.model.parameters()).dtype

    with torch.no_grad():
        if isinstance(X, torch.Tensor):
            X_tensor = X.to(self.device, dtype=dtype)
        else:
            X_tensor = torch.tensor(X, device=self.device, dtype=dtype)

        output = self.model(X_tensor)

    return output.detach().cpu().numpy()

__init__(input_model, device=None)

Wraps a PyTorch model to the superclass

Source code in interpretation/models/pytorch_model.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def __init__(
    self, 
    input_model: torch.nn.Module, 
    device: torch.device | str = None
) -> None:
    """Wraps a PyTorch model to the superclass"""

    if not isinstance(input_model, torch.nn.Module):
        raise TypeError("input_model must be an instance of torch.nn.Module")

    super().__init__(input_model)

    self.device = device or torch.device(
        "cuda" if torch.cuda.is_available() else "cpu"
    )
    self.model.to(self.device)
    self.model.eval()

    self.layer_dict = {
        name: module
        for name, module in self.model.named_modules()
        if len(list(module.children())) == 0
    }

activation_at_layer(X, layer_identifier)

Forward passes input X up to the specified layer. layer_identifier can be an integer index or layer name string.

Source code in interpretation/models/pytorch_model.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def activation_at_layer(
    self, 
    X: npt.NDArray | torch.Tensor, 
    layer_identifier: int | str | list[int | str]
) -> list[npt.NDArray]:
    """
    Forward passes input X up to the specified layer.
    layer_identifier can be an integer index or layer name string.
    """ 
    if not isinstance(layer_identifier, list):
        layer_identifier = [layer_identifier]

    target_modules = {}
    all_modules = [m for m in self.model.modules()][1:]

    for out_idx, identifier in enumerate(layer_identifier):
        if isinstance(identifier, int):
            target_module = all_modules[identifier]
        elif isinstance(identifier, str):
            target_module = dict(self.model.named_modules()).get(identifier)

        if target_module is None:
            raise ValueError(f"Layers '{layer_identifier}' not found in PyTorch model.")

        target_modules[target_module] = out_idx

    activations = [None] * len(layer_identifier)
    handles = []

    def hook_generator(out_idx):
        def hook(module, input, output):
            activations[out_idx] = output.detach().cpu().numpy()
        return hook

    for module, out_idx in target_modules.items():
        handles.append(module.register_forward_hook(hook_generator(out_idx))) #runs the hook code when forward passed to the layer

    try:
        self.__call__(X)
    finally:
        for h in handles:
            h.remove()

    return activations

compute_gradients(X, objective, objective_fn, wrt=None)

Computes the gradient \(\frac{\partial \text{Objective}}{\partial X}\).

objective is Objective, wrt is X.

The gradient is calculated after applying the objective function to the objective layer w.r.t the wrt layer.

Tuple forms of objective or wrt specifies specific neurons, e.g. (0, 1) means the second neuron of the first layer.

Note: When no wrt is provided, it will default to 'input'.

Source code in interpretation/models/pytorch_model.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def compute_gradients(
    self,
    X: npt.NDArray | torch.Tensor,
    objective: int | str | Tuple[int, int],
    objective_fn: Callable[[torch.Tensor], torch.tensor],
    wrt: int | str | Tuple[int, int] = None
) -> npt.NDArray:
    r"""Computes the gradient \(\frac{\partial \text{Objective}}{\partial X}\).

    `objective` is Objective, `wrt` is X.

    The gradient is calculated after applying the objective function to the objective layer w.r.t the wrt layer.

    Tuple forms of `objective` or `wrt` specifies specific neurons, e.g. (0, 1) means the second neuron of the first layer.

    Note: When no `wrt` is provided, it will default to `'input'`.
    """
    dtype = next(self.model.parameters()).dtype

    if isinstance(X, torch.Tensor):
        X_tensor = X.to(self.device, dtype=dtype)
    else:
        X_tensor = torch.tensor(X, device=self.device, dtype=dtype)        

    activations = {}
    hooks = []

    def hook_generator(name, neuron_idx=None):
        def hook(module, input, output):
            if neuron_idx is not None:
                activations[name] = output.squeeze()[neuron_idx]
            else:
                activations[name] = output
        return hook

    obj_name = self._layer_int_to_name(objective)
    wrt_name = self._layer_int_to_name(wrt) if wrt is not None else "input"

    if isinstance(objective, (int, str)):
        hooks.append(self.layer_dict[obj_name].register_forward_hook(hook_generator("obj")))
    else:
        hooks.append(self.layer_dict[obj_name].register_forward_hook(hook_generator("obj", objective[1])))

    if wrt_name != "input":
        hooks.append(self.layer_dict[wrt_name].register_forward_hook(hook_generator("wrt")))

    try:
        if wrt_name == "input":
            target_wrt = X_tensor
            target_wrt.requires_grad_()

            _ = self.model(target_wrt)
        else:
            def retain_grad_hook(module, input, output):
                output.retain_grad()
                activations["wrt"] = output

            hooks.append(self.layer_dict[wrt_name].register_forward_hook(retain_grad_hook))

            _ = self.model(X_tensor)
            target_wrt = activations["wrt"]

        obj_activation = activations["obj"]
        loss = objective_fn(obj_activation)

        if target_wrt.grad is not None:
            target_wrt.grad.zero_()
        loss.backward()

        grad = target_wrt.grad.detach().cpu().numpy()

        if isinstance(wrt, Tuple):
            return grad.squeeze()[wrt[1]]

        return grad

    finally:
        for h in hooks:
            h.remove()