Data Generation
x, y = load_data(n_points=1000, n_dims=10)
Next, we can use these data points to create a dataset and a data loader (no mini-batches this time):
Data Preparation
ball_dataset = TensorDataset(torch.as_tensor(x).float(), torch.as_tensor(y).float())
ball_loader = DataLoader(dataset=ball_dataset, batch_size=len(x))
torch.manual_seed(11) n_features = x.shape[1] n_layers = 5 n_hidden_units = 100 activation_fn = nn.ReLU model = build_model(n_features, n_layers, n_hidden_units, activation_fn, use_bn=False) model
Sequential( (h1): Linear(in_features=10, out_features=100, bias=True) (a1): ReLU() (h2): Linear(in_features=100, out_features=100, bias=True) (a2): ReLU() (h3): Linear(in_features=100, out_features=100, bias=True) (a3): ReLU() (h4): Linear(in_features=100, out_features=100, bias=True) (a4): ReLU() (h5): Linear(in_features=100, out_features=100, bias=True) (a5): ReLU() (o): Linear(in_features=100, out_features=1, bias=True) )
loss_fn = nn.BCEWithLogitsLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
hidden_layers = [f'h{i}' for i in range(1, n_layers + 1)] activation_layers = [f'a{i}' for i in range(1, n_layers + 1)] sbs = StepByStep(model, loss_fn, optimizer) sbs.set_loaders(ball_loader) sbs.capture_parameters(hidden_layers) sbs.capture_gradients(hidden_layers) sbs.attach_hooks(activation_layers) sbs.train(1)
parms, gradients, activations = get_plot_data(train_loader=ball_loader, model=model)
fig = plot_violins(parms, gradients, activations)
from torch.nn import functional as F, init class Linear(Module): r"""Applies an affine linear transformation to the incoming data: :math:`y = xA^T + b`. This module supports :ref:`TensorFloat32<tf32_on_ampere>`. On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward. Args: in_features: size of each input sample out_features: size of each output sample bias: If set to ``False``, the layer will not learn an additive bias. Default: ``True`` Shape: - Input: :math:`(*, H_{in})` where :math:`*` means any number of dimensions including none and :math:`H_{in} = \text{in\_features}`. - Output: :math:`(*, H_{out})` where all but the last dimension are the same shape as the input and :math:`H_{out} = \text{out\_features}`. Attributes: weight: the learnable weights of the module of shape :math:`(\text{out\_features}, \text{in\_features})`. The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`k = \frac{1}{\text{in\_features}}` bias: the learnable bias of the module of shape :math:`(\text{out\_features})`. If :attr:`bias` is ``True``, the values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where :math:`k = \frac{1}{\text{in\_features}}` Examples:: >>> m = nn.Linear(20, 30) >>> input = torch.randn(128, 20) >>> output = m(input) >>> print(output.size()) torch.Size([128, 30]) """ __constants__ = ["in_features", "out_features"] in_features: int out_features: int weight: Tensor def __init__( self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None, ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() self.in_features = in_features self.out_features = out_features self.weight = Parameter( torch.empty((out_features, in_features), **factory_kwargs) ) if bias: self.bias = Parameter(torch.empty(out_features, **factory_kwargs)) else: self.register_parameter("bias", None) self.reset_parameters() def reset_parameters(self) -> None: # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with # uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see # https://github.com/pytorch/pytorch/issues/57109 init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 init.uniform_(self.bias, -bound, bound) def forward(self, input: Tensor) -> Tensor: return F.linear(input, self.weight, self.bias) def extra_repr(self) -> str: return f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}"
Weight Initialization
def weights_init(model): if isinstance(model, nn.Linear): nn.init.kaiming_uniform_(model.weight, nonlinearity='relu') if model.bias is not None: nn.init.zeros_(model.bias)
with torch.no_grad():
model.apply(weights_init)
You should also use the no_grad() context manager while initializing / modifying the weights and biases of your model.
fig = plot_schemes(n_features, n_layers, n_hidden_units, ball_loader)
fig = plot_scheme_bn(n_features, n_layers, n_hidden_units, ball_loader)
from sklearn.datasets import make_regression x_reg, y_reg = make_regression(n_samples=1000, n_features=10, noise=0.1, random_state=42) x_reg = torch.as_tensor(x_reg).float() y_reg = torch.as_tensor(y_reg).float() dataset = TensorDataset(x_reg, y_reg) train_loader = DataLoader(dataset=dataset, batch_size=32, shuffle=True)
Even though we cannot plot a ten-dimensional regression, we can still visualize the distribution of both features and target values.
torch.manual_seed(11) model = nn.Sequential() model.add_module('fc1', nn.Linear(10, 15)) model.add_module('act', nn.ReLU()) model.add_module('fc2', nn.Linear(15, 1)) loss_fn = nn.MSELoss() optimizer = optim.SGD(model.parameters(), lr=0.01)
Before training it, let’s set our instance of the StepByStep class to capture the gradients for the weights in the hidden layer (fc1):
Model Training
sbs_reg = StepByStep(model, loss_fn, optimizer) sbs_reg.set_loaders(train_loader) sbs_reg.capture_gradients(['fc1'])
sbs_reg.train(2) sbs_reg.losses # [np.float64(17367.675262451172), np.float64(17306.57373046875)]
It turns out, two epochs are not enough to get exploding gradients.
Train another two epochs:
sbs_reg.train(2) sbs_reg.losses # [np.float64(17367.675262451172), # np.float64(17306.57373046875), # np.float64(17362.7119140625), # np.float64(17301.90805053711)]
Increase learning rate to get exploding gradients at the second epoch:
torch.manual_seed(11) model = nn.Sequential() model.add_module('fc1', nn.Linear(10, 15)) model.add_module('act', nn.ReLU()) model.add_module('fc2', nn.Linear(15, 1)) loss_fn = nn.MSELoss() optimizer = optim.SGD(model.parameters(), lr=0.03509966) sbs_reg = StepByStep(model, loss_fn, optimizer) sbs_reg.set_loaders(train_loader) sbs_reg.capture_gradients(['fc1']) sbs_reg.train(2) sbs_reg.losses # [np.float64(2.7670513733822185e+29), np.float64(nan)]
grads = np.array(sbs_reg._gradients['fc1']['weight']) print(grads.mean(axis=(1, 2)))
[ 2.44117905e-02 2.06106851e-01 1.81766347e-02 -2.35392176e+00
1.11135334e+01 -2.45726564e+01 -6.84081076e-01 -5.51755449e+00
2.17133981e+00 1.24698135e+00 -1.04324239e-01 9.37458067e-02
2.01123643e-01 5.32199112e-01 2.62268398e+01 -2.59524339e+01
6.88615245e-01 -1.83998085e-01 -3.75629093e+00 -1.86331088e-01
2.50075083e+00 1.52384045e-01 9.11065417e-03 6.15998185e-02
7.58073210e-01 4.96338612e-01 7.26849278e-01 7.53988466e-02
1.81392539e+01 -6.54686861e+01 1.38875571e+07 -2.59043075e+21
nan nan nan nan
nan nan nan nan
nan nan nan nan
nan nan nan nan
nan nan nan nan
nan nan nan nan
nan nan nan nan
nan nan nan nan]
The first NaN shows up at the 33rd update, but the explosion started at the 31st update: The average gradient goes from tens (1e+01) to ten millions (1e+07) in one step, to (1e+21) or whatever this is called in the next, to a full-blown NaN.
torch.manual_seed(42) parm = nn.Parameter(torch.randn(2, 1)) fake_grads = torch.tensor([[2.5], [.8]])
parm.grad = fake_grads.clone() # Gradient Value Clipping nn.utils.clip_grad_value_(parm, clip_value=1.0) parm.grad # tensor([[1.0000], # [0.8000]])
def clip_backprop(model, clip_value): handles = [] for p in model.parameters(): if p.requires_grad: func = lambda grad: torch.clamp(grad, -clip_value, clip_value) handle = p.register_hook(func) handles.append(handle) return handles
Do not forget that you should remove the hooks using handle.remove() after you’re done with them.
parm.grad = fake_grads.clone() # Gradient Norm Clipping nn.utils.clip_grad_norm_(parm, max_norm=1.0, norm_type=2) print(fake_grads) print(fake_grads.norm()) print(parm.grad) print(parm.grad.norm())
tensor([[2.5000], [0.8000]]) tensor(2.6249) tensor([[0.9524], [0.3048]]) tensor(1.0000)
self.clipping = None def set_clip_grad_value(self, clip_value): self.clipping = lambda: nn.utils.clip_grad_value_(self.model.parameters(), clip_value=clip_value) def set_clip_grad_norm(self, max_norm, norm_type=2): self.clipping = lambda: nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=max_norm, norm_type=norm_type) def remove_clip(self): self.clipping = None def _make_train_step_fn(self): # Build function that performs a step in the train loop def perform_train_step_fn(x, y): self.model.train() yhat = self.model(x) loss = self.loss_fn(yhat, y) loss.backward() if callable(self.clipping): self.clipping() self.optimizer.step() self.optimizer.zero_grad() return loss.item() return perform_train_step_fn
def set_clip_backprop(self, clip_value): if self.clipping is None: self.clipping = [] for p in self.model.parameters(): if p.requires_grad: func = lambda grad: torch.clamp(grad, -clip_value, clip_value) handle = p.register_hook(func) self.clipping.append(handle) def remove_clip(self): if isinstance(self.clipping, list): for handle in self.clipping: handle.remove() self.clipping = None
Let’s use the weights_init function to initialize the weights so we can reset the parameters of our model that had its gradients exploded.
Moreover, let’s use a ten times higher learning rate; after all, we’re in full control of the gradients now:
Model Configuration
torch.manual_seed(42)
with torch.no_grad():
model.apply(weights_init)
optimizer = optim.SGD(model.parameters(), lr=0.03509966)
Before training it, let’s use set_clip_grad_value() to make sure no gradients are ever above 1.0:
Model Training
sbs_reg_clip = StepByStep(model, loss_fn, optimizer) sbs_reg_clip.set_loaders(train_loader) sbs_reg_clip.set_clip_grad_value(1.0) sbs_reg_clip.capture_gradients(['fc1']) sbs_reg_clip.train(10) sbs_reg_clip.remove_clip() sbs_reg_clip.remove_hooks()
avg_grad = np.array(sbs_reg_clip._gradients['fc1']['weight']).mean(axis=(1, 2)) avg_grad.min(), avg_grad.max() # (np.float64(-6.1622835213318465), np.float64(1.5875422400461199))
torch.manual_seed(42)
with torch.no_grad():
model.apply(weights_init)
sbs_reg_clip_hook = StepByStep(model, loss_fn, optimizer) sbs_reg_clip_hook.set_loaders(train_loader) sbs_reg_clip_hook.set_clip_backprop(1.0) sbs_reg_clip_hook.capture_gradients(['fc1']) sbs_reg_clip_hook.train(10) sbs_reg_clip_hook.remove_clip() sbs_reg_clip_hook.remove_hooks()
torch.manual_seed(42) with torch.no_grad(): model.apply(weights_init) optimizer = optim.SGD(model.parameters(), lr=0.01) sbs_reg_clip_hook = StepByStep(model, loss_fn, optimizer) sbs_reg_clip_hook.set_loaders(train_loader) sbs_reg_clip_hook.set_clip_backprop(1.0) sbs_reg_clip_hook.capture_gradients(['fc1']) sbs_reg_clip_hook.train(10) sbs_reg_clip_hook.remove_clip() sbs_reg_clip_hook.remove_hooks()
