fix: use torch.amp via compatibility shim to avoid deprecated torch.cuda.amp warnings

funasr/bin/train.py

@@ -16,7 +16,7 @@
 import torch.distributed as dist
 
 from omegaconf import DictConfig, OmegaConf
-from torch.cuda.amp import autocast, GradScaler
+from funasr.utils.amp import autocast, GradScaler
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.algorithms.join import Join

funasr/bin/train_ds.py

@@ -15,7 +15,7 @@
 import torch.distributed as dist
 
 from omegaconf import DictConfig, OmegaConf
-from torch.cuda.amp import autocast, GradScaler
+from funasr.utils.amp import autocast, GradScaler
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.algorithms.join import Join

funasr/models/bat/model.py

@@ -24,7 +24,7 @@
 
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/bicif_paraformer/model.py

@@ -26,7 +26,7 @@
 from funasr.train_utils.device_funcs import to_device
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/campplus/model.py

@@ -26,7 +26,7 @@
 
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/contextual_paraformer/model.py

@@ -29,7 +29,7 @@
 
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/ct_transformer/model.py

@@ -23,7 +23,7 @@
 except:
     pass
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/ct_transformer_streaming/model.py

@@ -16,7 +16,7 @@
 
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/data2vec/data2vec.py

@@ -22,7 +22,7 @@
 from funasr.train_utils.device_funcs import force_gatherable
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/e_paraformer/model.py

@@ -8,7 +8,7 @@
 import copy
 import torch
 import logging
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 from typing import Union, Dict, List, Tuple, Optional
 
 from funasr.register import tables

funasr/models/e_paraformer/pif_predictor.py

@@ -11,7 +11,7 @@
 from funasr.register import tables
 from funasr.train_utils.device_funcs import to_device
 from funasr.models.transformer.utils.nets_utils import make_pad_mask
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 
 
 @tables.register("predictor_classes", "PifPredictor")

funasr/models/emotion2vec/model.py

@@ -23,7 +23,7 @@
 
 logger = logging.getLogger(__name__)
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/fsmn_kws/model.py

@@ -6,7 +6,7 @@
 import time
 import torch
 import logging
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 from typing import Union, Dict, List, Tuple, Optional
 
 from funasr.register import tables

funasr/models/fsmn_kws_mt/model.py

@@ -6,7 +6,7 @@
 import time
 import torch
 import logging
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 from typing import Union, Dict, List, Tuple, Optional
 
 from funasr.register import tables

funasr/models/lcbnet/model.py

@@ -9,7 +9,7 @@
 import time
 import torch
 import torch.nn as nn
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 
 from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
 from funasr.models.ctc.ctc import CTC

funasr/models/llm_asr/model.py

@@ -5,7 +5,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 import re
 from funasr.models.scama.utils import sequence_mask
 from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
@@ -583,7 +583,7 @@ def forward(
 
         batch_size, frames, _ = speech.shape
 
-        with torch.cuda.amp.autocast(enabled=False):
+        with autocast(enabled=False):
             # audio encoder
             encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
 
@@ -618,7 +618,7 @@ def forward(
                     batch_idx, :min_len, :
                 ]
 
-        with torch.cuda.amp.autocast(
+        with autocast(
             enabled=True if self.llm_dtype != "fp32" else False, dtype=dtype_map[self.llm_dtype]
         ):
             labels_ids[labels_ids == -1] = -100
@@ -883,7 +883,7 @@ def inference(
             llm_dtype = "fp16" if kwargs.get("fp16", False) else llm_dtype
             llm_dtype = "bf16" if kwargs.get("bf16", False) else llm_dtype
 
-        with torch.cuda.amp.autocast(
+        with autocast(
             enabled=True if llm_dtype != "fp32" else False, dtype=dtype_map[llm_dtype]
         ):
             label = contents["assistant"][0]
@@ -1159,7 +1159,7 @@ def forward(
         batch_size_speech, frames, _ = speech.shape
         batch_size, token_num = input_ids.shape
 
-        with torch.cuda.amp.autocast(enabled=False):
+        with autocast(enabled=False):
             # audio encoder
             encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
 
@@ -1203,7 +1203,7 @@ def forward(
 
                     speech_idx += 1
 
-        with torch.cuda.amp.autocast(
+        with autocast(
             enabled=True if self.llm_dtype != "fp32" else False, dtype=dtype_map[self.llm_dtype]
         ):
             labels_ids[labels_ids == -1] = -100
@@ -1551,7 +1551,7 @@ def inference(
             llm_dtype = "fp16" if kwargs.get("fp16", False) else llm_dtype
             llm_dtype = "bf16" if kwargs.get("bf16", False) else llm_dtype
 
-        with torch.cuda.amp.autocast(
+        with autocast(
             enabled=True if llm_dtype != "fp32" else False, dtype=dtype_map[llm_dtype]
         ):
             label = contents["assistant"][-1]

funasr/models/llm_asr_nar/model.py

@@ -5,7 +5,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 
 from funasr.models.scama.utils import sequence_mask
 from funasr.losses.label_smoothing_loss import LabelSmoothingLoss

funasr/models/mfcca/e2e_asr_mfcca.py

@@ -25,7 +25,7 @@
 from funasr.models.base_model import FunASRModel
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/monotonic_aligner/model.py

@@ -6,7 +6,7 @@
 import time
 import copy
 import torch
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 from typing import Union, Dict, List, Tuple, Optional
 
 from funasr.register import tables

funasr/models/paraformer/cif_predictor.py

@@ -10,7 +10,7 @@
 from funasr.register import tables
 from funasr.train_utils.device_funcs import to_device
 from funasr.models.transformer.utils.nets_utils import make_pad_mask
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 
 
 @tables.register("predictor_classes", "CifPredictor")

funasr/models/paraformer/model.py

@@ -7,7 +7,7 @@
 import copy
 import torch
 import logging
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 from typing import Union, Dict, List, Tuple, Optional
 
 from funasr.register import tables

funasr/models/paraformer_streaming/model.py

@@ -26,7 +26,7 @@
 
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/paraformer_v2_community/model.py

@@ -7,7 +7,7 @@
 import copy
 import torch
 import logging
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 from typing import Union, Dict, List, Tuple, Optional
 
 from funasr.register import tables

funasr/models/sa_asr/e2e_sa_asr.py

@@ -29,7 +29,7 @@
 from funasr.models.base_model import FunASRModel
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/sanm_kws/model.py

@@ -6,7 +6,7 @@
 import time
 import torch
 import logging
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 from typing import Union, Dict, List, Tuple, Optional
 
 from funasr.register import tables

funasr/models/sanm_kws_streaming/model.py

@@ -26,7 +26,7 @@
 
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/scama/model.py

@@ -28,7 +28,7 @@
 from funasr.models.scama.utils import sequence_mask
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/seaco_paraformer/model.py

@@ -33,7 +33,7 @@
 
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/sense_voice/model.py

@@ -6,7 +6,7 @@
 import torch.nn.functional as F
 from torch import Tensor
 from torch import nn
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 from funasr.metrics.compute_acc import compute_accuracy, th_accuracy
 from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
 from funasr.train_utils.device_funcs import force_gatherable

funasr/models/sond/e2e_diar_sond.py

@@ -29,7 +29,7 @@
 from funasr.utils.hinter import hint_once
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/transducer/model.py

@@ -23,7 +23,7 @@
 
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/models/transformer/model.py

@@ -4,7 +4,7 @@
 import time
 import torch
 import torch.nn as nn
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 
 from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
 from funasr.models.ctc.ctc import CTC

funasr/models/uniasr/model.py

@@ -6,7 +6,7 @@
 import time
 import torch
 import logging
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 from typing import Union, Dict, List, Tuple, Optional
 
 from funasr.register import tables

funasr/models/whisper_lid/model.py

@@ -5,7 +5,7 @@
 import torch
 import numpy as np
 import torch.nn as nn
-from torch.cuda.amp import autocast
+from funasr.utils.amp import autocast
 
 from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
 from funasr.models.ctc.ctc import CTC

funasr/models/xvector/e2e_sv.py

@@ -31,7 +31,7 @@
 from funasr.models.base_model import FunASRModel
 
 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
-    from torch.cuda.amp import autocast
+    from funasr.utils.amp import autocast
 else:
     # Nothing to do if torch<1.6.0
     @contextmanager

funasr/train_utils/trainer.py

@@ -6,7 +6,7 @@
 from tqdm import tqdm
 from datetime import datetime
 import torch.distributed as dist
-from torch.cuda.amp import autocast, GradScaler
+from funasr.utils.amp import autocast, GradScaler
 from contextlib import nullcontext, contextmanager
 from pathlib import Path
 

funasr/train_utils/trainer_ds.py

@@ -6,7 +6,7 @@
 from tqdm import tqdm
 from datetime import datetime
 import torch.distributed as dist
-from torch.cuda.amp import autocast, GradScaler
+from funasr.utils.amp import autocast, GradScaler
 from contextlib import nullcontext, contextmanager
 from pathlib import Path
 from torch.nn.parallel import DistributedDataParallel as DDP
@@ -32,7 +32,7 @@ def maybe_autocast(dtype=None, use_deepspeed=False):
             use_deepspeed: TODO.
         """
     if use_deepspeed:
-        with torch.cuda.amp.autocast(enabled=True, dtype=dtype, cache_enabled=False):
+        with autocast(enabled=True, dtype=dtype, cache_enabled=False):
             yield
     else:
         if dtype == torch.float16 or dtype == torch.bfloat16: