|
""" |
|
Overview: |
|
In this file, we provide a set of utility functions for probing network parameters and gradients, |
|
which can be helpful in analyzing and debugging the inner workings of various models. |
|
""" |
|
from typing import List, Tuple |
|
import numpy as np |
|
import torch |
|
import torch.nn as nn |
|
|
|
|
|
def renormalize(inputs: torch.Tensor, first_dim: int = 1) -> torch.Tensor: |
|
""" |
|
Overview: |
|
Normalize the input data using the max-min-normalization. |
|
Arguments: |
|
- inputs (:obj:`torch.Tensor`): The input data needs to be normalized. |
|
- first_dim (:obj:`int`): The first dimension of flattening the input data. |
|
Returns: |
|
- output (:obj:`torch.Tensor`): The normalized data. |
|
""" |
|
if first_dim < 0: |
|
first_dim = len(inputs.shape) + first_dim |
|
flat_input = inputs.view(*inputs.shape[:first_dim], -1) |
|
max_val = torch.max(flat_input, first_dim, keepdim=True).values |
|
min_val = torch.min(flat_input, first_dim, keepdim=True).values |
|
flat_input = (flat_input - min_val) / (max_val - min_val) |
|
|
|
return flat_input.view(*input.shape) |
|
|
|
|
|
def get_dynamic_mean(model: nn.Module) -> float: |
|
dynamic_mean = np.abs(model.conv.weight.detach().cpu().numpy().reshape(-1)).tolist() |
|
|
|
for block in model.resblocks: |
|
for name, param in block.named_parameters(): |
|
dynamic_mean += np.abs(param.detach().cpu().numpy().reshape(-1)).tolist() |
|
dynamic_mean = sum(dynamic_mean) / len(dynamic_mean) |
|
return dynamic_mean |
|
|
|
|
|
def get_reward_mean(model: nn.Module) -> Tuple[np.ndarray, float]: |
|
reward_w_dist = model.conv1x1_reward.weight.detach().cpu().numpy().reshape(-1) |
|
|
|
for name, param in model.fc.named_parameters(): |
|
temp_weights = param.detach().cpu().numpy().reshape(-1) |
|
reward_w_dist = np.concatenate((reward_w_dist, temp_weights)) |
|
reward_mean = np.abs(reward_w_dist).mean() |
|
return reward_w_dist, reward_mean |
|
|
|
|
|
def get_params_mean(model: nn.Module) -> Tuple[np.ndarray, float, float, float]: |
|
representation_mean = model.representation_network.get_param_mean() |
|
dynamic_mean = model.dynamics_network.get_dynamic_mean() |
|
reward_w_dist, reward_mean = model.dynamics_network.get_reward_mean() |
|
|
|
return reward_w_dist, representation_mean, dynamic_mean, reward_mean |
|
|
|
|
|
def get_gradients(model: nn.Module) -> List[torch.Tensor]: |
|
grads = [] |
|
for p in model.parameters(): |
|
grad = None if p.grad is None else p.grad.detach() |
|
grads.append(grad) |
|
return grads |
|
|
|
|
|
def set_gradients(model: nn.Module, gradients: List[torch.Tensor]) -> None: |
|
|
|
for g, p in zip(gradients, model.parameters()): |
|
if g is not None: |
|
p.grad = g |
|
|
