Spaces:
Configuration error
Configuration error
File size: 1,611 Bytes
4efbc62 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
# International Conference on Computer Vision (ICCV), 2023
import copy
import math
import torch
import torch.nn as nn
from efficientvit.models.utils import is_parallel
__all__ = ["EMA"]
def update_ema(
ema: nn.Module, new_state_dict: dict[str, torch.Tensor], decay: float
) -> None:
for k, v in ema.state_dict().items():
if v.dtype.is_floating_point:
v -= (1.0 - decay) * (v - new_state_dict[k].detach())
class EMA:
def __init__(self, model: nn.Module, decay: float, warmup_steps=2000):
self.shadows = copy.deepcopy(
model.module if is_parallel(model) else model
).eval()
self.decay = decay
self.warmup_steps = warmup_steps
for p in self.shadows.parameters():
p.requires_grad = False
def step(self, model: nn.Module, global_step: int) -> None:
with torch.no_grad():
msd = (model.module if is_parallel(model) else model).state_dict()
update_ema(
self.shadows,
msd,
self.decay * (1 - math.exp(-global_step / self.warmup_steps)),
)
def state_dict(self) -> dict[float, dict[str, torch.Tensor]]:
return {self.decay: self.shadows.state_dict()}
def load_state_dict(self, state_dict: dict[float, dict[str, torch.Tensor]]) -> None:
for decay in state_dict:
if decay == self.decay:
self.shadows.load_state_dict(state_dict[decay])
|