FishLeg package

Submodules

FishLeg.fishleg module

class FishLeg.fishleg.FishLeg(model: Module, draw: Callable[[Module, Tensor], Tuple[Tensor, Tensor]], nll: Callable[[Module, Tuple[Tensor, Tensor]], Tensor], aux_dataloader: DataLoader, likelihood: Optional[FishLikelihood] = None, fish_lr: float = 0.05, damping: float = 0.5, weight_decay: float = 1e-05, beta: float = 0.9, update_aux_every: int = 10, aux_lr: float = 0.0001, aux_betas: Tuple[float, float] = (0.9, 0.999), aux_eps: float = 1e-08, num_steps=None, batch_speedup: bool = False, full: bool = True, normalization: bool = False, fine_tune: bool = False, module_names: List[str] = [], skip_names: List[str] = [], initialization: str = 'uniform', scale: float = 1.0, warmup: int = 0, warmup_data: Optional[DataLoader] = None, warmup_loss: Optional[Callable] = None, device: str = 'cpu', config=None, verbose=False)

Bases: Optimizer

Implement FishLeg algorithm.

As described in https://openreview.net/forum?id=c9lAOPvQHS.

Parameters

• model (torch.nn.Module) – a pytorch neural network module, can be nested in a tree structure

• draw (Callable[[nn.Module, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]) – Sampling function that takes a model \(f\) and input data \(\mathbf X\), and returns \((\mathbf X, \mathbf y)\), where \(\mathbf y\) is sampled from the conditional distribution \(p(\mathbf y|f(\mathbf X))\)

• nll (Callable[[nn.Module, Tuple[torch.Tensor, torch.Tensor]], torch.Tensor]) – A function that takes a model and data, and evaluate the negative log-likelihood.

• aux_dataloader (torch.utiles.data.DataLoader) – A function that takes a batch size as input and output dataset with corresponding size.

:param FishLikelihood likelihooda FishLeg likelihood, with Qv method if

any parameters are learnable.

Parameters

• fish_lr (float) – Learning rate, for the parameters of the input model using FishLeg (default: 1e-2)

• damping (float) – Static damping applied to Fisher matrix, \(\gamma\), for stability when FIM becomes near-singular. (default: 5e-1)

• weight_decay (float) – L2 penalty on weights (default: 1e-5)

• beta (float) – coefficient for running averages of gradient (default: 0.9)

• update_aux_every (int) – Number of iteration after which an auxiliary update is executed, if negative, then run -update_aux_every auxiliary updates in each outer iteration. (default: 10)

• aux_lr (float) – learning rate for the auxiliary parameters, using Adam (default: 1e-3)

• aux_betas (Tuple[float, float]) – Coefficients used for computing running averages of gradient and its square for auxiliary parameters (default: (0.9, 0.999))

• aux_eps (float) – Term added to the denominator to improve numerical stability for auxiliary parameters (default: 1e-8)

• batch_speedup (bool) – Whether to use speed-up Qv product (default: False)

• full (bool) – Whether to use full inner and outer diagonal rescalling for block Kronecker approximation of Q. (default: True)

• normalization (bool) – Whether to use normalization on gradients when calculating the auxiliary loss, this is important to learn about curvature even when gradients are small (default: False)

• fine_tune (bool) – Whether to use Fisher as preconditioner of pretrained tasks, and fine-tune on a downstream task. If True, Q will be fixed and continual learning will be performed (default: False)

• module_names (List) – A List of module names wished to be optimized/pruned by FishLeg. (default: [], meaning all modules optimized/pruned by FishLeg)

• initialization (string) – Initialization of weights (default: uniform)

• warmup (int) – If warmup is zero, the default SGD warmup will be used, where Q is initialized as a scaled identity matrix. If warmup is positive, the diagonal of Q will be initialized as \(\frac{1}{g^2 + \gamma}\); and in this case, warmup_data and warmup_loss should be provided for sampling of gradients.

• scale (float) – Help specify initial scale of the inverse Fisher Information matrix approximation. If using SGD warmup we suggest, \(\eta=\gamma^{-1}\). If warmup is positive, scale should be 1. (default: 1)

• device (str) – The device where calculations will be performed using PyTorch Tensors.

Example:

>>> aux_loader = torch.utils.data.DataLoader(train_data, shuffle=True, batch_size=100)
>>> train_loader = torch.utils.data.DataLoader(train_data, shuffle=True, batch_size=100)
>>>
>>> likelihood = FixedGaussianLikelihood(sigma=1.0)
>>>
>>> def nll(model, data_x, data_y):
>>>     pred_y = model.forward(data_x)
>>>     return likelihood.nll(data_y, pred_y)
>>>
>>> def draw(model, data_x):
>>>     pred_y = model.forward(data_x)
>>>     return likelihood.draw(pred_y)
>>>
>>> model = nn.Sequential(
>>>     nn.Linear(2, 5),
>>>     nn.ReLU(),
>>>     nn.Linear(5, 1),
>>> )
>>>
>>> opt = FishLeg(
>>>     model,
>>>     draw,
>>>     nll,
>>>     aux_loader
>>> )
>>>
>>> for data_x, data_y in dataloader:
>>>     opt.zero_grad()
>>>     pred_y = model(data_x)
>>>     loss = nn.MSELoss()(data_y, pred_y)
>>>     loss.backward()
>>>     opt.step()
>>>     if iteration % 10 == 0:
>>>         print(loss.detach())

init_model_aux(model: Module, module_names: List[str], skip_names: List[str], config=None) → Union[Module, List]

Given a model to optimize, parameters can be devided to

1. those fixed as pre-trained.

2. those required to optimize using FishLeg.

Replace modules in the second group with FishLeg modules.

Args:

model (torch.nn.Module, required):

A model containing modules to replace with FishLeg modules containing extra functionality related to FishLeg algorithm.

Returns:

torch.nn.Module, the replaced model.

pretrain_fish(dataloader: DataLoader, loss: Callable[[Module, Tuple[Tensor, Tensor]], Tensor], iterations: int = 10000, difference: bool = False, verbose: bool = False, testloader: Optional[DataLoader] = None, batch_size: int = 500, fisher: bool = True) → List

step(closure=None) → None

Performes a single optimization step of FishLeg.

update_aux(train=True, fisher=True) → None

Performs a single auxliarary parameter update using Adam. By minimizing the following objective:

\[nll(model, \theta + \epsilon Q(\lambda)g) + nll(model, \theta - \epsilon Q(\lambda)g) - 2\epsilon^2g^T Q(\lambda)g\]

where \(\theta\) is the parameters of model, \(\lambda\) is the auxliarary parameters.

warmup_aux(dataloader: Optional[DataLoader] = None, loss: Optional[Callable[[Module, Tuple[Tensor, Tensor]], Tensor]] = None, scale: float = 1.0) → None

Warm up auxilirary parameters, if warmup is larger zero, follow approxiamte Adam, if warmup is zero, follow SGD

FishLeg.fishleg_layers module

class FishLeg.fishleg_layers.FishBatchNorm2d(num_features: int, eps: float = 1e-05, momentum: float = 0.1, affine: bool = True, track_running_stats: bool = True, init_scale=1.0, device=None, dtype=None)

Bases: BatchNorm2d, FishModule

Qv(v: Tuple, full=False)

\(Q(\lambda)\) is a positive definite matrix which will effectively estimate the inverse damped Fisher Information Matrix. Appropriate choices for \(Q\) should take into account the architecture of the model/module. It is usually parameterized as a positive definite Kronecker-factored block-diagonal matrix, with block sizes reflecting the layer structure of the neural networks.

Args:

aux: (Dict, required): auxiliary parameters,

\(\lambda\), a dictionary with keys, the name of the auxiliary parameters, and values, the auxiliary parameters of the module. These auxiliaray parameters will form \(Q(\lambda)\).

v: (Tuple[Tensor, …], required): Values of the original parameters,

in an order that align with self.order, to multiply with \(Q(\lambda)\).

full: (bool, optional), whether to use full inner and outer re-scaling

Returns:

Tuple[Tensor, …]: The calculated \(Q(\lambda)v\) products,

in same order with self.order.

affine: bool

diagQ()

eps: float

momentum: float

num_features: int

track_running_stats: bool

class FishLeg.fishleg_layers.FishConv2d(in_channels: int, out_channels: int, kernel_size, stride=1, padding=0, dilation=1, groups: int = 1, bias: bool = True, padding_mode: str = 'zeros', device=None)

Bases: Conv2d, FishModule

Qv(v: Tuple[Tensor, Optional[Tensor]], full: bool = False) → Tuple[Tensor, Optional[Tensor]]

Inspired by KFAC’s conv2D layer by Grosse and Martens: Kronecker product of sizes (out_channels ⊗ (in_channels_eff * k_size))

bias: Optional[Tensor]

diagQ() → Tensor

Similar maths as the Linear layer

dilation: Tuple[int, ...]

groups: int

in_channels: int

kernel_size: Tuple[int, ...]

out_channels: int

output_padding: Tuple[int, ...]

padding: Union[str, Tuple[int, ...]]

padding_mode: str

stride: Tuple[int, ...]

transposed: bool

warmup(v: Optional[Tuple[Tensor, Tensor]] = None, init_scale: float = 1.0) → None

weight: Tensor

class FishLeg.fishleg_layers.FishLayerNorm(normalized_shape, eps: float = 1e-05, elementwise_affine: bool = True, init_scale=1.0, device=None, dtype=None)

Bases: LayerNorm, FishModule

Qv(v: Tuple, full=False)

\(Q(\lambda)\) is a positive definite matrix which will effectively estimate the inverse damped Fisher Information Matrix. Appropriate choices for \(Q\) should take into account the architecture of the model/module. It is usually parameterized as a positive definite Kronecker-factored block-diagonal matrix, with block sizes reflecting the layer structure of the neural networks.

Args:

aux: (Dict, required): auxiliary parameters,

\(\lambda\), a dictionary with keys, the name of the auxiliary parameters, and values, the auxiliary parameters of the module. These auxiliaray parameters will form \(Q(\lambda)\).

v: (Tuple[Tensor, …], required): Values of the original parameters,

in an order that align with self.order, to multiply with \(Q(\lambda)\).

full: (bool, optional), whether to use full inner and outer re-scaling

Returns:

Tuple[Tensor, …]: The calculated \(Q(\lambda)v\) products,

in same order with self.order.

diagQ()

elementwise_affine: bool

eps: float

normalized_shape: Tuple[int, ...]

class FishLeg.fishleg_layers.FishLinear(in_features: int, out_features: int, bias: bool = True, device=None, dtype=None)

Bases: Linear, FishModule

Qg() → Tuple[Tensor, Tensor]

Speed up Qg product, when batch size is smaller than parameter size. By chain rule:

\[DW_i = g_i\hat{a}^T_{i-1}\]

where \(DW_i\) is gradient of parameter of the ith layer, \(g_i\) is gradient w.r.t output of ith layer and \(\hat{a}_i\) is input to ith layer, and output of (i-1)th layer.

Qv(v: Tuple[Tensor, Tensor], full: bool = False) → Tuple[Tensor, Tensor]

For fully-connected layers, the default structure of \(Q\) as a block-diaglonal matrix is, .. math:

Q_l = (R_lR_l^T \otimes L_lL_l^T)

where \(l\) denotes the l-th layer. The matrix \(R_l\) has size \((N_{l-1} + 1) \times (N_{l-1} + 1)\) while the matrix \(L_l\) has size \(N_l \times N_l\). The auxiliarary parameters \(\lambda\) are represented by the matrices \(L_l, R_l\). For a Kronecker form that introduces full inner and outer diagonal rescaling structure is,

\[Q_l = A_l(L_l \otimes R_l^T) D_l^2 (L_l^T \otimes R_l) A_l\]

where \(A_l\) and \(D_l\) are two diagonal matrices of the appropriate size.

diagQ() → Tuple

The Q matrix defines the inverse fisher approximation as below:

\[Q_l = (R_lR_l^T \otimes L_lL_l^T)\]

where \(l\) denotes the l-th layer. The matrix \(R_l\) has size \((N_{l-1} + 1) \times (N_{l-1} + 1)\) while the matrix \(L_l\) has size \(N_l \times N_l\). The auxiliarary parameters \(\lambda\) are represented by the matrices \(L_l, R_l\).

The diagonal of this matrix is therefore calculated by

\[\text{diag}(Q_l) = \text{diag}(R_l R_l^T) \otimes \text{diag}(L_l L_l^T)\]

where \(\text{diag}\) involves summing over the columns of the and \(\otimes\) remains as the Kronecker product.

in_features: int

out_features: int

save_layer_grad_output(grad_output: Tuple[Tensor, ...]) → None

save_layer_input(input_: List[Tensor]) → None

warmup(v: Optional[Tuple[Tensor, Tensor]] = None, batch_speedup: bool = False, init_scale: float = 1.0) → None

weight: Tensor

FishLeg.fishleg_likelihood module

class FishLeg.fishleg_likelihood.BernoulliLikelihood(device: str = 'cpu')

Bases: FishLikelihood

The Bernoulli likelihood used for classification. Using the standard Normal CDF \(\Phi(x)\)) and the identity \(\Phi(-x) = 1-\Phi(x)\), we can write the likelihood as:

\[p(y|f(x))=\Phi(yf(x))\]

draw(preds: Tensor) → Tensor

Draw samples from the conditional distribution \(p(\mathbf y|f(\mathbf x))\)

Parameters

preds (torch.Tensor) – Predictions from model \(f(\mathbf x)\)

nll(preds: Tensor, observations: Tensor) → Tensor

Computes the negative log-likelihood \(\ell(\theta, \mathcal D)=-\log p(\mathbf y|f(\mathbf x))\)

Parameters

• observations (torch.Tensor) – Values of \(y\).

• preds (torch.Tensor) – Predictions from model \(f(\mathbf x)\)

Return type

torch.Tensor

class FishLeg.fishleg_likelihood.FishLikelihood

Bases: object

A Likelihood in FishLeg specifies a probablistic modeling, which attributes the mapping from latent function values \(f(\mathbf X)\) to observed labels \(y\).

For example, in the case of regression, a Gaussian likelihood can be chosen, as

\[y(\mathbf x) = f(\mathbf x) + \epsilon, \:\:\:\: \epsilon \sim N(0,\sigma^{2}_{n} \mathbf I)\]

As for the case of classification, a Bernoulli distribution can be chosen

\[\begin{split}y(\mathbf x) = \begin{cases} 1 & \text{w/ probability} \:\: \sigma(f(\mathbf x)) \\ 0 & \text{w/ probability} \:\: 1-\sigma(f(\mathbf x)) \end{cases}\end{split}\]

abstract draw(preds, **kwargs)

Draw samples from the conditional distribution \(p(\mathbf y|f(\mathbf x))\)

Parameters

preds (torch.Tensor) – Predictions from model \(f(\mathbf x)\)

get_parameters() → List

return a list of learnable parameter.

abstract nll(preds, observations, **kwargs)

Computes the negative log-likelihood \(\ell(\theta, \mathcal D)=-\log p(\mathbf y|f(\mathbf x))\)

Parameters

• observations (torch.Tensor) – Values of \(y\).

• preds (torch.Tensor) – Predictions from model \(f(\mathbf x)\)

Return type

torch.Tensor

class FishLeg.fishleg_likelihood.FixedGaussianLikelihood(sigma: Tensor, device: str = 'cpu')

Bases: FishLikelihood

The standard likelihood for regression, but assuming fixed heteroscedastic noise.

\[p(y | f(x)) = f(x) + \epsilon, \:\:\:\: \epsilon \sim N(0,\sigma^{2})\]

Parameters

sigma (torch.Tensor) – Known observation standard deviation for each example.

draw(preds: Tensor) → Tensor

Draw samples from the conditional distribution \(p(\mathbf y|f(\mathbf x))\)

Parameters

preds (torch.Tensor) – Predictions from model \(f(\mathbf x)\)

property get_variance: Tensor

nll(preds: Tensor, observations: Tensor) → Tensor

Computes the negative log-likelihood \(\ell(\theta, \mathcal D)=-\log p(\mathbf y|f(\mathbf x))\)

Parameters

• observations (torch.Tensor) – Values of \(y\).

• preds (torch.Tensor) – Predictions from model \(f(\mathbf x)\)

Return type

torch.Tensor

class FishLeg.fishleg_likelihood.GaussianLikelihood(sigma: Tensor, device: str = 'cpu')

Bases: FishLikelihood

The standard likelihood for regression, but assuming fixed heteroscedastic noise.

\[p(y | f(x)) = f(x) + \epsilon, \:\:\:\: \epsilon \sim N(0,\sigma^{2})\]

Parameters

sigma (torch.Tensor) – standard deviation for each example; also to be learned during training.

Qv(v) → List

draw(preds: Tensor) → Tensor

Draw samples from the conditional distribution \(p(\mathbf y|f(\mathbf x))\)

Parameters

preds (torch.Tensor) – Predictions from model \(f(\mathbf x)\)

get_aux_parameters() → List

get_parameters() → List

return a list of learnable parameter.

init_aux(init_scale) → None

nll(preds: Tensor, observations: Tensor) → Tensor

Computes the negative log-likelihood \(\ell(\theta, \mathcal D)=-\log p(\mathbf y|f(\mathbf x))\)

Parameters

• observations (torch.Tensor) – Values of \(y\).

• preds (torch.Tensor) – Predictions from model \(f(\mathbf x)\)

Return type

torch.Tensor

class FishLeg.fishleg_likelihood.SoftMaxLikelihood(device: str = 'cpu')

Bases: FishLikelihood

draw(preds: Tensor) → Tensor

Draw samples from the conditional distribution \(p(\mathbf y|f(\mathbf x))\)

Parameters

preds (torch.Tensor) – Predictions from model \(f(\mathbf x)\)

nll(preds: Tensor, observations: Tensor) → Tensor

Computes the negative log-likelihood \(\ell(\theta, \mathcal D)=-\log p(\mathbf y|f(\mathbf x))\)

Parameters

• observations (torch.Tensor) – Values of \(y\).

• preds (torch.Tensor) – Predictions from model \(f(\mathbf x)\)

Return type

torch.Tensor

FishLeg.utils module

FishLeg.utils.get_named_layers_by_regex(module: Module, param_names: List[str], params_strict: bool = False) → List[NamedLayerParam]

Parameters

• module – the module to get the matching layers and params from

• param_names – a list of names or regex patterns to match with full parameter paths. Regex patterns must be specified with the prefix ‘re:’

• params_strict – if True, this function will raise an exception if there a parameter is not found to match every name or regex in param_names

Returns

a list of NamedLayerParam tuples whose full parameter names in the given module match one of the given regex patterns or parameter names

FishLeg.utils.recursive_getattr(obj, attr)

FishLeg.utils.recursive_setattr(obj, attr, value)

FishLeg.utils.update_dict(replace: Module, module: Module) → Module

Module contents