Spaces:
Sleeping
Sleeping
File size: 2,525 Bytes
9043dc9 |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
# 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 efficientvit.models.efficientvit import (
EfficientViTSeg,
efficientvit_seg_b0,
efficientvit_seg_b1,
efficientvit_seg_b2,
efficientvit_seg_b3,
efficientvit_seg_l1,
efficientvit_seg_l2,
)
from efficientvit.models.nn.norm import set_norm_eps
from efficientvit.models.utils import load_state_dict_from_file
__all__ = ["create_seg_model"]
REGISTERED_SEG_MODEL: dict[str, dict[str, str]] = {
"cityscapes": {
"b0": "assets/checkpoints/seg/cityscapes/b0.pt",
"b1": "assets/checkpoints/seg/cityscapes/b1.pt",
"b2": "assets/checkpoints/seg/cityscapes/b2.pt",
"b3": "assets/checkpoints/seg/cityscapes/b3.pt",
################################################
"l1": "assets/checkpoints/seg/cityscapes/l1.pt",
"l2": "assets/checkpoints/seg/cityscapes/l2.pt",
},
"ade20k": {
"b1": "assets/checkpoints/seg/ade20k/b1.pt",
"b2": "assets/checkpoints/seg/ade20k/b2.pt",
"b3": "assets/checkpoints/seg/ade20k/b3.pt",
################################################
"l1": "assets/checkpoints/seg/ade20k/l1.pt",
"l2": "assets/checkpoints/seg/ade20k/l2.pt",
},
}
def create_seg_model(
name: str, dataset: str, pretrained=True, weight_url: str or None = None, **kwargs
) -> EfficientViTSeg:
model_dict = {
"b0": efficientvit_seg_b0,
"b1": efficientvit_seg_b1,
"b2": efficientvit_seg_b2,
"b3": efficientvit_seg_b3,
#########################
"l1": efficientvit_seg_l1,
"l2": efficientvit_seg_l2,
}
model_id = name.split("-")[0]
if model_id not in model_dict:
raise ValueError(f"Do not find {name} in the model zoo. List of models: {list(model_dict.keys())}")
else:
model = model_dict[model_id](dataset=dataset, **kwargs)
if model_id in ["l1", "l2"]:
set_norm_eps(model, 1e-7)
if pretrained:
weight_url = weight_url or REGISTERED_SEG_MODEL[dataset].get(name, None)
if weight_url is None:
raise ValueError(f"Do not find the pretrained weight of {name}.")
else:
weight = load_state_dict_from_file(weight_url)
model.load_state_dict(weight)
return model
|