mirror of
https://github.com/modelscope/FunASR.git
synced 2026-10-12 02:19:18 +00:00
谢翊凡
9413d7d768
Use the non-deprecated torch.amp APIs when available while preserving the torch.cuda.amp fallback and legacy autocast argument semantics on older PyTorch releases.
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
import importlib.util
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from funasr.utils.amp import autocast
|
|
|
|
|
|
@pytest.mark.parametrize("args", [(False,), (False, torch.float16, False)])
|
|
def test_autocast_preserves_legacy_positional_arguments(args):
|
|
with autocast(*args):
|
|
pass
|
|
|
|
|
|
def test_amp_import_falls_back_when_torch_amp_is_missing(monkeypatch):
|
|
legacy_autocast = object()
|
|
legacy_grad_scaler = object()
|
|
torch_stub = types.ModuleType("torch")
|
|
torch_stub.__path__ = []
|
|
cuda_stub = types.ModuleType("torch.cuda")
|
|
cuda_stub.__path__ = []
|
|
cuda_amp_stub = types.ModuleType("torch.cuda.amp")
|
|
cuda_amp_stub.autocast = legacy_autocast
|
|
cuda_amp_stub.GradScaler = legacy_grad_scaler
|
|
torch_stub.cuda = cuda_stub
|
|
cuda_stub.amp = cuda_amp_stub
|
|
|
|
monkeypatch.setitem(sys.modules, "torch", torch_stub)
|
|
monkeypatch.setitem(sys.modules, "torch.cuda", cuda_stub)
|
|
monkeypatch.setitem(sys.modules, "torch.cuda.amp", cuda_amp_stub)
|
|
|
|
module_path = Path(__file__).parents[1] / "funasr" / "utils" / "amp.py"
|
|
spec = importlib.util.spec_from_file_location("funasr_amp_legacy_test", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
|
|
assert module.autocast is legacy_autocast
|
|
assert module.GradScaler is legacy_grad_scaler
|