mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Add SimpleNorm to create_norm factory

Browse Source
This commit is contained in:
Ross Wightman 2024-12-30 14:22:42 -08:00
parent 5809c2fe5e
commit a4146b79d1
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
timm/layers/create_norm.py
View File

@ -10,7 +10,7 @@ from typing import Type
import torch.nn as nn import torch.nn as nn
from .norm import GroupNorm, GroupNorm1, LayerNorm, LayerNorm2d, RmsNorm, RmsNorm2d from .norm import GroupNorm, GroupNorm1, LayerNorm, LayerNorm2d, RmsNorm, RmsNorm2d, SimpleNorm, SimpleNorm2d
from torchvision.ops.misc import FrozenBatchNorm2d from torchvision.ops.misc import FrozenBatchNorm2d
_NORM_MAP = dict( _NORM_MAP = dict(
@ -23,6 +23,8 @@ _NORM_MAP = dict(
layernorm2d=LayerNorm2d, layernorm2d=LayerNorm2d,
rmsnorm=RmsNorm, rmsnorm=RmsNorm,
rmsnorm2d=RmsNorm2d, rmsnorm2d=RmsNorm2d,
simplenorm=SimpleNorm,
simplenorm2d=SimpleNorm2d,
frozenbatchnorm2d=FrozenBatchNorm2d, frozenbatchnorm2d=FrozenBatchNorm2d,
) )
_NORM_TYPES = {m for n, m in _NORM_MAP.items()} _NORM_TYPES = {m for n, m in _NORM_MAP.items()}