Skip to content

Utilization

get_supported_optimizers(filters=None)

Return list of available optimizer names, sorted alphabetically.

Parameters:

Name Type Description Default
filters Optional[Union[str, List[str]]]

Optional[Union[str, List[str]]]. wildcard filter string that works with fmatch. if None, it will return the whole list.

 None
Source code in pytorch_optimizer/optimizer/__init__.py 
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
def get_supported_optimizers(filters: Optional[Union[str, List[str]]] = None) -> List[str]:
    r"""Return list of available optimizer names, sorted alphabetically.

    :param filters: Optional[Union[str, List[str]]]. wildcard filter string that works with fmatch.
        if None, it will return the whole list.
    """
    if filters is None:
        return sorted(OPTIMIZERS.keys())

    include_filters: Sequence[str] = filters if isinstance(filters, (tuple, list)) else [filters]

    filtered_list: Set[str] = set()
    for include_filter in include_filters:
        filtered_list.update(fnmatch.filter(OPTIMIZERS.keys(), include_filter))

    return sorted(filtered_list)

get_supported_lr_schedulers(filters=None)

Return list of available lr scheduler names, sorted alphabetically.

Parameters:

Name Type Description Default
filters Optional[Union[str, List[str]]]

Optional[Union[str, List[str]]]. wildcard filter string that works with fmatch. if None, it will return the whole list.

 None
Source code in pytorch_optimizer/lr_scheduler/__init__.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def get_supported_lr_schedulers(filters: Optional[Union[str, List[str]]] = None) -> List[str]:
    r"""Return list of available lr scheduler names, sorted alphabetically.

    :param filters: Optional[Union[str, List[str]]]. wildcard filter string that works with fmatch. if None, it will
        return the whole list.
    """
    if filters is None:
        return sorted(LR_SCHEDULERS.keys())

    include_filters: Sequence[str] = filters if isinstance(filters, (tuple, list)) else [filters]

    filtered_list: Set[str] = set()
    for include_filter in include_filters:
        filtered_list.update(fnmatch.filter(LR_SCHEDULERS.keys(), include_filter))

    return sorted(filtered_list)

get_supported_loss_functions(filters=None)

Return list of available loss function names, sorted alphabetically.

Parameters:

Name Type Description Default
filters Optional[Union[str, List[str]]]

Optional[Union[str, List[str]]]. wildcard filter string that works with fmatch. if None, it will return the whole list.

 None
Source code in pytorch_optimizer/loss/__init__.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def get_supported_loss_functions(filters: Optional[Union[str, List[str]]] = None) -> List[str]:
    r"""Return list of available loss function names, sorted alphabetically.

    :param filters: Optional[Union[str, List[str]]]. wildcard filter string that works with fmatch. if None, it will
        return the whole list.
    """
    if filters is None:
        return sorted(LOSS_FUNCTIONS.keys())

    include_filters: Sequence[str] = filters if isinstance(filters, (tuple, list)) else [filters]

    filtered_list: Set[str] = set()
    for include_filter in include_filters:
        filtered_list.update(fnmatch.filter(LOSS_FUNCTIONS.keys(), include_filter))

    return sorted(filtered_list)

CPUOffloadOptimizer

Offload optimizer to CPU for single-GPU training. This will reduce GPU memory by the size of optimizer state.

Reference: https://github.com/pytorch/ao/blob/main/torchao/prototype/low_bit_optim/cpu_offload.py

Parameters:

Name Type Description Default
params PARAMETERS

PARAMETERS. a list of parameters or parameter groups.

 required
optimizer_class Type[Optimizer]

Type[torch.optim.Optimizer]. constructor of the base optimizer. Defaults to :class:torch.optim.AdamW.

 AdamW
offload_gradients bool

bool. free GPU gradients once they are moved to CPU. Not compatible with gradient accumulation.

 False
kwargs

other keyword arguments to be passed to the base optimizer e.g. lr, weight_decay.

 {}
Source code in pytorch_optimizer/optimizer/utils.py 
 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
class CPUOffloadOptimizer:  # pragma: no cover
    """Offload optimizer to CPU for single-GPU training. This will reduce GPU memory by the size of optimizer state.

    Reference: https://github.com/pytorch/ao/blob/main/torchao/prototype/low_bit_optim/cpu_offload.py

    :param params: PARAMETERS. a list of parameters or parameter groups.
    :param optimizer_class: Type[torch.optim.Optimizer]. constructor of the base optimizer. Defaults to
        :class:`torch.optim.AdamW`.
    :param offload_gradients: bool. free GPU gradients once they are moved to CPU. Not compatible with gradient
        accumulation.
    :param kwargs: other keyword arguments to be passed to the base optimizer e.g. `lr`, `weight_decay`.
    """

    def __init__(
        self,
        params: PARAMETERS,
        optimizer_class: Type[Optimizer] = torch.optim.AdamW,
        *,
        offload_gradients: bool = False,
        **kwargs,
    ) -> None:
        if optimizer_class is torch.optim.AdamW and TORCH_VERSION_AT_LEAST_2_4 and 'fused' not in kwargs:
            kwargs.update(fused=True)

        param_groups = list(params)
        if len(param_groups) == 0:
            raise ValueError('optimizer got an empty parameter list')

        if not isinstance(param_groups[0], dict):
            param_groups = [{'params': param_groups}]

        self.param_cuda2cpu_map: Dict[torch.Tensor, torch.Tensor] = {}
        self.optim_dict: Dict[torch.Tensor, Optimizer] = {}
        self.stream = torch.cuda.Stream()

        self.queue = {}

        def backward_hook(p_cuda: torch.Tensor) -> None:
            if p_cuda.grad is None:
                return

            p_cpu = self.param_cuda2cpu_map[p_cuda]

            self.stream.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(self.stream):
                p_cpu.grad.copy_(p_cuda.grad, non_blocking=True)

            if p_cuda in self.queue:
                del self.queue[p_cuda]

            self.queue[p_cuda] = self.stream.record_event()

            if offload_gradients:
                p_cuda.grad.record_stream(self.stream)
                p_cuda.grad = None

        for param_group in param_groups:
            params = param_group.pop('params')
            if params is None:
                continue

            for p_cuda in params:
                p_cpu = torch.empty_like(p_cuda, device='cpu', pin_memory=True)
                p_cpu.grad = torch.empty_like(p_cpu, pin_memory=True)

                p_cpu.copy_(p_cuda.detach(), non_blocking=True)
                self.param_cuda2cpu_map[p_cuda] = p_cpu

                p_cuda.register_post_accumulate_grad_hook(backward_hook)
                self.optim_dict[p_cuda] = optimizer_class([{'params': p_cpu, **param_group}], **kwargs)  # type: ignore

    @torch.no_grad()
    def step(self, closure: CLOSURE = None) -> LOSS:
        loss = None
        if closure is not None:
            loss = closure()

        for p_cuda, grad_d2h_event in self.queue.items():
            grad_d2h_event.synchronize()
            self.optim_dict[p_cuda].step()

            p_cpu = self.param_cuda2cpu_map[p_cuda]
            with torch.cuda.stream(self.stream):
                p_cuda.copy_(p_cpu, non_blocking=True)

        self.queue.clear()

        return loss

    def zero_grad(self, _: bool = True) -> None:
        for p_cuda in self.param_cuda2cpu_map:
            p_cuda.grad = None

    @property
    def param_groups(self):
        return functools.reduce(operator.add, (optim.param_groups for optim in self.optim_dict.values()), [])

    def state_dict(self):
        return [optim.state_dict() for optim in self.optim_dict.values()]

    def load_state_dict(self, state_dict):
        for optim, optim_state_dict in zip(self.optim_dict.values(), state_dict):
            optim.load_state_dict(optim_state_dict)

is_valid_parameters(parameters)

Check where the parameters are valid.

Source code in pytorch_optimizer/optimizer/utils.py 
213
214
215
def is_valid_parameters(parameters: PARAMETERS) -> bool:
    r"""Check where the parameters are valid."""
    return isinstance(parameters, (list, tuple)) and len(parameters) > 0 and isinstance(parameters[0], dict)

has_overflow(grad_norm)

Detect inf and NaN in grad_norm.

Source code in pytorch_optimizer/optimizer/utils.py 
218
219
220
def has_overflow(grad_norm: torch.Tensor) -> bool:
    r"""Detect inf and NaN in grad_norm."""
    return bool(torch.logical_or(torch.isnan(grad_norm), torch.isinf(grad_norm)).any())

to_real(x)

Return real value of tensor.

Source code in pytorch_optimizer/optimizer/utils.py 
223
224
225
def to_real(x: torch.Tensor) -> torch.Tensor:
    r"""Return real value of tensor."""
    return x.real if torch.is_complex(x) else x

normalize_gradient(x, use_channels=False, epsilon=1e-08)

Normalize gradient with stddev.

Parameters:

Name Type Description Default
x Tensor

torch.Tensor. gradient.

 required
use_channels bool

bool. channel-wise normalization.

 False
epsilon float

float. eps.

 1e-08
Source code in pytorch_optimizer/optimizer/utils.py 
228
229
230
231
232
233
234
235
236
237
238
239
240
241
def normalize_gradient(x: torch.Tensor, use_channels: bool = False, epsilon: float = 1e-8) -> None:
    r"""Normalize gradient with stddev.

    :param x: torch.Tensor. gradient.
    :param use_channels: bool. channel-wise normalization.
    :param epsilon: float. eps.
    """
    size: int = x.dim()
    if size > 1 and use_channels:
        s = x.std(dim=tuple(range(1, size)), keepdim=True).add_(epsilon)
        x.div_(s)
    elif torch.numel(x) > 2:
        s = x.std().add_(epsilon)
        x.div_(s)

clip_grad_norm(parameters, max_norm=0.0, sync=False)

Clip gradient norms.

During combination with FSDP, will also ensure that grad norms are aggregated across all workers,
since each worker only stores their shard of the gradients.

Parameters:

Name Type Description Default
parameters Union[PARAMETERS, Tensor]

PARAMETERS. Parameters whose gradients we wish to clip.

 required
max_norm float

float. Maximum norm we wish the gradients to have. If non-positive, then we will not perform clipping.

 0.0
sync bool

bool. Boolean indicating whether we should aggregate across the distributed group. Used only in combination with FSDP.

 False

Returns:

Type Description
Union[Tensor, float]

The gradient norm across all parameters, before clipping.
Source code in pytorch_optimizer/optimizer/utils.py 
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
def clip_grad_norm(
    parameters: Union[PARAMETERS, torch.Tensor],
    max_norm: float = 0.0,
    sync: bool = False,
) -> Union[torch.Tensor, float]:
    r"""Clip gradient norms.

        During combination with FSDP, will also ensure that grad norms are aggregated across all workers,
        since each worker only stores their shard of the gradients.

    :param parameters: PARAMETERS. Parameters whose gradients we wish to clip.
    :param max_norm: float. Maximum norm we wish the gradients to have. If non-positive, then we will not perform
        clipping.
    :param sync: bool. Boolean indicating whether we should aggregate across the distributed group. Used only in
        combination with FSDP.
    :returns: The gradient norm across all parameters, before clipping.
    """
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]

    # make sure any generators are expanded
    parameters = list(parameters)

    # if syncing we need to manually perform the clipping so that we aggregate properly
    if max_norm > 0 and not sync:
        return clip_grad_norm_(parameters, max_norm)

    norm_sq = sum(p.grad.norm() ** 2 for p in parameters if p.grad is not None)
    if sync:  # pragma: no cover
        # also need to get the norms from all the other sharded works in FSDP
        all_reduce(norm_sq)

    grad_norm: float = math.sqrt(norm_sq)
    if max_norm > 0:  # pragma: no cover
        clip_coefficient = max_norm / (grad_norm + 1e-6)
        for p in parameters:
            if p.grad is not None:
                p.grad.detach().mul_(clip_coefficient)

    return grad_norm

unit_norm(x, norm=2.0)

Get norm of unit.

Source code in pytorch_optimizer/optimizer/utils.py 
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
def unit_norm(x: torch.Tensor, norm: float = 2.0) -> torch.Tensor:
    r"""Get norm of unit."""
    keep_dim: bool = True
    dim: Optional[Union[int, Tuple[int, ...]]] = None

    x_len: int = len(x.shape)
    if x_len <= 1:
        keep_dim = False
    elif x_len in (2, 3):
        dim = 1
    elif x_len == 4:
        dim = (1, 2, 3)
    else:
        dim = tuple(range(1, x_len))

    return x.norm(p=norm, dim=dim, keepdim=keep_dim)

disable_running_stats(model)

Disable running stats (momentum) of BatchNorm.

Source code in pytorch_optimizer/optimizer/utils.py 
304
305
306
307
308
309
310
311
312
def disable_running_stats(model: nn.Module):
    r"""Disable running stats (momentum) of BatchNorm."""

    def _disable(module):
        if isinstance(module, _BatchNorm):
            module.backup_momentum = module.momentum
            module.momentum = 0

    model.apply(_disable)

enable_running_stats(model)

Enable running stats (momentum) of BatchNorm.

Source code in pytorch_optimizer/optimizer/utils.py 
315
316
317
318
319
320
321
322
def enable_running_stats(model: nn.Module):
    r"""Enable running stats (momentum) of BatchNorm."""

    def _enable(module):
        if isinstance(module, _BatchNorm) and hasattr(module, 'backup_momentum'):
            module.momentum = module.backup_momentum

    model.apply(_enable)

get_global_gradient_norm(param_groups)

Get global gradient norm.

Source code in pytorch_optimizer/optimizer/utils.py 
325
326
327
328
329
330
331
332
333
334
335
@torch.no_grad()
def get_global_gradient_norm(param_groups: List[Dict]) -> torch.Tensor:
    r"""Get global gradient norm."""
    global_grad_norm = torch.zeros(1, dtype=torch.float32, device=param_groups[0]['params'][0].device)

    for group in param_groups:
        for p in group['params']:
            if p.grad is not None:
                global_grad_norm.add_(p.grad.norm().pow(2))

    return global_grad_norm

reg_noise(network1, network2, num_data, lr, eta=0.008, temperature=0.0001)

Entropy-MCMC: Sampling from flat basins with ease.

usage: https://github.com/lblaoke/EMCMC/blob/master/exp/cifar10_emcmc.py

Parameters:

Name Type Description Default
network1 Module

nn.Module. network.

 required
network2 Module

nn.Module. network.

 required
num_data int

int. number of training data.

 required
lr float

float. learning rate.

 required
eta float

float. eta.

 0.008
temperature float

float. temperature.

 0.0001
Source code in pytorch_optimizer/optimizer/utils.py 
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
@torch.no_grad()
def reg_noise(
    network1: nn.Module, network2: nn.Module, num_data: int, lr: float, eta: float = 8e-3, temperature: float = 1e-4
) -> Union[torch.Tensor, float]:
    r"""Entropy-MCMC: Sampling from flat basins with ease.

    usage: https://github.com/lblaoke/EMCMC/blob/master/exp/cifar10_emcmc.py

    :param network1: nn.Module. network.
    :param network2: nn.Module. network.
    :param num_data: int. number of training data.
    :param lr: float. learning rate.
    :param eta: float. eta.
    :param temperature: float. temperature.
    """
    reg_coef: float = 0.5 / (eta * num_data)
    noise_coef: float = math.sqrt(2.0 / lr / num_data * temperature)

    loss = torch.tensor(0.0, device=next(network1.parameters()).device)

    for param1, param2 in zip(network1.parameters(), network2.parameters()):
        reg = (param1 - param2).pow_(2).mul_(reg_coef).sum()

        noise = param1 * torch.randn_like(param1)
        noise.add_(param2 * torch.randn_like(param2))

        loss.add_(reg - noise.mul_(noise_coef).sum())

    return loss

copy_stochastic(target, source)

Copy stochastic.

reference: https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905

Parameters:

Name Type Description Default
target Tensor

torch.Tensor. bfloat16 tensor.

 required
source Tensor

torch.Tensor. float32 tensor.

 required
Source code in pytorch_optimizer/optimizer/utils.py 
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
@torch.no_grad()
def copy_stochastic(target: torch.Tensor, source: torch.Tensor) -> None:
    r"""Copy stochastic.

    reference: https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905

    :param target: torch.Tensor. bfloat16 tensor.
    :param source: torch.Tensor. float32 tensor.
    """
    result = torch.randint_like(
        source,
        dtype=torch.int32,
        low=0,
        high=1 << 16,
    )

    result.add_(source.view(dtype=torch.int32))

    result.bitwise_and_(-65536)

    target.copy_(result.view(dtype=torch.float32))