Spaces:
Configuration error
Configuration error
File size: 818 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 |
# 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
from functools import partial
import torch.nn as nn
from efficientvit.models.utils import build_kwargs_from_config
__all__ = ["build_act"]
# register activation function here
REGISTERED_ACT_DICT: dict[str, type] = {
"relu": nn.ReLU,
"relu6": nn.ReLU6,
"hswish": nn.Hardswish,
"silu": nn.SiLU,
"gelu": partial(nn.GELU, approximate="tanh"),
}
def build_act(name: str, **kwargs) -> nn.Module or None:
if name in REGISTERED_ACT_DICT:
act_cls = REGISTERED_ACT_DICT[name]
args = build_kwargs_from_config(kwargs, act_cls)
return act_cls(**args)
else:
return None
|