towhee.models.layers.sam.SAM

class towhee.models.layers.sam.SAM(params, base_optimizer, rho=0.05, adaptive=False, **kwargs)

Bases: Optimizer

Sharpness-Aware Minimization that simultaneously minimizes loss value and loss sharpness. In particular, it seeks parameters that lie in neighborhoods having uniformly low loss.

Methods

add_param_group	Add a param group to the Optimizer s param_groups.
first_step
load_state_dict	Loads the optimizer state.
profile_hook_step
register_step_post_hook	Register an optimizer step post hook which will be called after optimizer step. It should have the following signature::.
register_step_pre_hook	Register an optimizer step pre hook which will be called before optimizer step. It should have the following signature::.
second_step
state_dict	Returns the state of the optimizer as a dict.
step	Performs a single optimization step (parameter update).
zero_grad	Sets the gradients of all optimized torch.Tensor s to zero.

__init__(params, base_optimizer, rho=0.05, adaptive=False, **kwargs)

__repr__()

Return repr(self).

add_param_group(param_group)

Add a param group to the Optimizer s param_groups.

This can be useful when fine tuning a pre-trained network as frozen layers can be made trainable and added to the Optimizer as training progresses.

Parameters:

param_group (dict) – Specifies what Tensors should be optimized along with group specific optimization options.

load_state_dict(state_dict)

Loads the optimizer state.

Parameters:

state_dict (dict) – optimizer state. Should be an object returned from a call to state_dict().

register_step_post_hook(hook: Callable[[...], None]) → RemovableHandle

Register an optimizer step post hook which will be called after optimizer step. It should have the following signature:

hook(optimizer, args, kwargs) -> None

The optimizer argument is the optimizer instance being used.

Parameters:

hook (Callable) – The user defined hook to be registered.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemoveableHandle

register_step_pre_hook(hook: Callable[[...], None]) → RemovableHandle

Register an optimizer step pre hook which will be called before optimizer step. It should have the following signature:

hook(optimizer, args, kwargs) -> None or modified args and kwargs

The optimizer argument is the optimizer instance being used. If args and kwargs are modified by the pre-hook, then the transformed values are returned as a tuple containing the new_args and new_kwargs.

Parameters:

hook (Callable) – The user defined hook to be registered.

Returns:

a handle that can be used to remove the added hook by calling handle.remove()

Return type:

torch.utils.hooks.RemoveableHandle

state_dict()

Returns the state of the optimizer as a dict.

It contains two entries:

• state - a dict holding current optimization state. Its content

  differs between optimizer classes.

• param_groups - a list containing all parameter groups where each

  parameter group is a dict

step(closure=None)

Performs a single optimization step (parameter update).

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

zero_grad(set_to_none: bool = True)

Sets the gradients of all optimized torch.Tensor s to zero.

Parameters:

set_to_none (bool) – instead of setting to zero, set the grads to None. This will in general have lower memory footprint, and can modestly improve performance. However, it changes certain behaviors. For example: 1. When the user tries to access a gradient and perform manual ops on it, a None attribute or a Tensor full of 0s will behave differently. 2. If the user requests zero_grad(set_to_none=True) followed by a backward pass, .grads are guaranteed to be None for params that did not receive a gradient. 3. torch.optim optimizers have a different behavior if the gradient is 0 or None (in one case it does the step with a gradient of 0 and in the other it skips the step altogether).