18 lines
719 B
Python
18 lines
719 B
Python
from .optimizer import Optimizer
|
|
from ..nn.modules import Module
|
|
from .lr_scheduler import _LRScheduler
|
|
from .. import device, Tensor
|
|
from typing import Iterable, Any, Optional, Callable, Union, List
|
|
|
|
class AveragedModel(Module):
|
|
def __init__(self, model: Module, device: Union[int, device]=...,
|
|
avg_fn: Callable[[Tensor, Tensor, int], Tensor]=...) -> None:...
|
|
|
|
def update_parameters(self, model: Module) -> None:...
|
|
|
|
def update_bn(loader: Iterable[Any], model: Module, device: Union[int, device]=...) -> None:...
|
|
|
|
class SWALR(_LRScheduler):
|
|
def __init__(self, optimizer: Optimizer, swa_lr: float, anneal_epochs: int,
|
|
anneal_strategy: str, last_epoch: int=...) -> None:...
|