mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
eb7653614f
* refactor activations into basic PyTorch, jit scripted, and memory efficient custom auto * implement hard-mish, better grad for hard-swish * add initial VovNet V1/V2 impl, fix #151 * VovNet and DenseNet first models to use NormAct layers (support BatchNormAct2d, EvoNorm, InplaceIABN) * Wrap IABN for any models that use it * make more models torchscript compatible (DPN, PNasNet, Res2Net, SelecSLS) and add tests
75 lines
1.6 KiB
Python
75 lines
1.6 KiB
Python
""" Model / Layer Config Singleton
|
|
"""
|
|
from typing import Any
|
|
|
|
__all__ = ['is_exportable', 'is_scriptable', 'set_exportable', 'set_scriptable', 'is_no_jit', 'set_no_jit']
|
|
|
|
# Set to True if prefer to have layers with no jit optimization (includes activations)
|
|
_NO_JIT = False
|
|
|
|
# Set to True if prefer to have activation layers with no jit optimization
|
|
_NO_ACTIVATION_JIT = False
|
|
|
|
# Set to True if exporting a model with Same padding via ONNX
|
|
_EXPORTABLE = False
|
|
|
|
# Set to True if wanting to use torch.jit.script on a model
|
|
_SCRIPTABLE = False
|
|
|
|
|
|
def is_no_jit():
|
|
return _NO_JIT
|
|
|
|
|
|
class set_no_jit:
|
|
def __init__(self, mode: bool) -> None:
|
|
global _NO_JIT
|
|
self.prev = _NO_JIT
|
|
_NO_JIT = mode
|
|
|
|
def __enter__(self) -> None:
|
|
pass
|
|
|
|
def __exit__(self, *args: Any) -> bool:
|
|
global _NO_JIT
|
|
_NO_JIT = self.prev
|
|
return False
|
|
|
|
|
|
def is_exportable():
|
|
return _EXPORTABLE
|
|
|
|
|
|
class set_exportable:
|
|
def __init__(self, mode: bool) -> None:
|
|
global _EXPORTABLE
|
|
self.prev = _EXPORTABLE
|
|
_EXPORTABLE = mode
|
|
|
|
def __enter__(self) -> None:
|
|
pass
|
|
|
|
def __exit__(self, *args: Any) -> bool:
|
|
global _EXPORTABLE
|
|
_EXPORTABLE = self.prev
|
|
return False
|
|
|
|
|
|
def is_scriptable():
|
|
return _SCRIPTABLE
|
|
|
|
|
|
class set_scriptable:
|
|
def __init__(self, mode: bool) -> None:
|
|
global _SCRIPTABLE
|
|
self.prev = _SCRIPTABLE
|
|
_SCRIPTABLE = mode
|
|
|
|
def __enter__(self) -> None:
|
|
pass
|
|
|
|
def __exit__(self, *args: Any) -> bool:
|
|
global _SCRIPTABLE
|
|
_SCRIPTABLE = self.prev
|
|
return False
|