mirror of
https://github.com/open-mmlab/mmpretrain.git
synced 2025-06-03 14:59:18 +08:00
1a7cebe4b9
* Refactor unit tests folder structure. * Remove label smooth and Vit test in `test_classifiers.py` * Rename test_utils in dataset to test_dataset_utils * Split test_models/test_utils/test_utils.py to multiple sub files. * Add unit tests of classifiers and heads * Use patch context manager. * Add unit test of `is_tracing`, and add warning in `is_tracing` if torch verison is smaller than 1.6.0
42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
import collections.abc
|
|
import warnings
|
|
from distutils.version import LooseVersion
|
|
from itertools import repeat
|
|
|
|
import torch
|
|
|
|
|
|
def is_tracing() -> bool:
|
|
if LooseVersion(torch.__version__) >= LooseVersion('1.6.0'):
|
|
on_trace = torch.jit.is_tracing()
|
|
# In PyTorch 1.6, torch.jit.is_tracing has a bug.
|
|
# Refers to https://github.com/pytorch/pytorch/issues/42448
|
|
if isinstance(on_trace, bool):
|
|
return on_trace
|
|
else:
|
|
return torch._C._is_tracing()
|
|
else:
|
|
warnings.warn(
|
|
'torch.jit.is_tracing is only supported after v1.6.0. '
|
|
'Therefore is_tracing returns False automatically. Please '
|
|
'set on_trace manually if you are using trace.', UserWarning)
|
|
return False
|
|
|
|
|
|
# From PyTorch internals
|
|
def _ntuple(n):
|
|
|
|
def parse(x):
|
|
if isinstance(x, collections.abc.Iterable):
|
|
return x
|
|
return tuple(repeat(x, n))
|
|
|
|
return parse
|
|
|
|
|
|
to_1tuple = _ntuple(1)
|
|
to_2tuple = _ntuple(2)
|
|
to_3tuple = _ntuple(3)
|
|
to_4tuple = _ntuple(4)
|
|
to_ntuple = _ntuple
|