update for adapter

training/benchmarks/wav2vec2/pytorch/README.md → training/benchmarks/wav2vec2/README.md

@@ -9,6 +9,7 @@ https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/SpeechRecogni
 
 ### 数据集下载地址(global proxy)
 http://www.openslr.org/resources/12
+请参考 https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechRecognition/wav2vec2#quick-start-guide 下载和处理数据
 
 
 ### 框架与芯片支持情况

training/benchmarks/wav2vec2/pytorch/config/model_config.py

@@ -14,7 +14,7 @@
 encoder_attention_heads = 12
 feature_grad_mult = 0.1
 ema = 0.0
-optimizer = "adam"
+optimizer = "fused_adam"
 clip_norm = 25
 weight_decay = 0.01
 lr_policy = "poly"

training/benchmarks/wav2vec2/pytorch/optimizer/__init__.py

@@ -1,17 +0,0 @@
-
-from common.fairseq.optim.fused_adam import get_fused_adam_class
-from common.utils import print_once
-
-def create_optimizer(model, args):
-
-    kw = {'lr': args.lr, 'weight_decay': args.weight_decay}
-    if args.optimizer == 'adam' and not (args.fp16 or args.bf16):
-        print_once('WARNING: Using FusedAdam instead of Adam')
-        kw.update({'betas': args.adam_betas, 'eps': args.adam_eps})
-        fused_adam_cls = get_fused_adam_class()
-        print(fused_adam_cls, "fused_adam_cls")
-        optimizer = fused_adam_cls(model.parameters(), **kw)
-    else:
-        raise ValueError(f'Invalid optimizer "{args.optimizer}"')
-
-    return optimizer

training/benchmarks/wav2vec2/pytorch/train/trainer.py

@@ -22,7 +22,6 @@
 from model import create_model
 from train.evaluator import Evaluator
 from train.training_state import TrainingState
-from optimizer import create_optimizer
 from loss.criterion import Wav2vecCriterion
 
 
@@ -50,7 +49,7 @@ def __init__(self, driver: Driver, adapter, evaluator: Evaluator,
     def init(self):
         self.model = create_model(self.config)
         self.model = self.init_model(self.model, self.device)
-        self.optimizer = create_optimizer(self.model, self.config)
+        self.optimizer = self.adapter.create_optimizer(self.model, self.config)
         self.optim = self.optimizer
 
         Metrics = W2v2Metrics
@@ -153,7 +152,7 @@ def train_one_epoch(self, config, epoch, step, train_dataloader, sampler):
                 grads_mult_factor = world_size / self.metrics.partials[
                     'sample_size']
 
-                if self.config.optimizer == 'adam' and not (self.config.fp16
+                if self.config.optimizer == 'fused_adam' and not (self.config.fp16
                                                        or self.config.bf16):
                     # adam and non-amp optimizer - can use 'scale' kwarg for step
                     # and defer grad multiplication
@@ -168,7 +167,7 @@ def train_one_epoch(self, config, epoch, step, train_dataloader, sampler):
                         # calculate grad norm, maybe clip
                         grad_norm = self.optim.clip_grad_norm(self.config.clip_norm)
 
-                    if self.config.optimizer == 'adam' and not (self.config.fp16
+                    if self.config.optimizer == 'fused_adam' and not (self.config.fp16
                                                            or self.config.bf16):
                         self.scaler.step(self.optim,
                                          scale=1. / grads_mult_factor)

training/benchmarks/wav2vec2/pytorch/train/trainer_adapter.py

@@ -1,10 +1,11 @@
 import torch.distributed as dist
 import config
 
-from torch import nn, Tensor
-from driver.dist_pytorch import main_proc_print
+from torch import nn
 from torch.nn.parallel import DistributedDataParallel as DDP
-
+
+from common.fairseq.optim.fused_adam import get_fused_adam_class
+from driver.dist_pytorch import main_proc_print
 
 
 def convert_model(model: nn.Module) -> nn.Module:
@@ -27,3 +28,17 @@ def model_to_ddp(model: nn.Module) -> nn.Module:
         from common.fairseq.dist import ModuleProxyWrapper
         model = ModuleProxyWrapper(model)
     return model
+
+
+def create_optimizer(model, args):
+
+    kw = {'lr': args.lr, 'weight_decay': args.weight_decay}
+    if args.optimizer == 'fused_adam' and not (args.fp16 or args.bf16):
+        kw.update({'betas': args.adam_betas, 'eps': args.adam_eps})
+        fused_adam_cls = get_fused_adam_class()
+        print(fused_adam_cls, "fused_adam_cls")
+        optimizer = fused_adam_cls(model.parameters(), **kw)
+    else:
+        raise ValueError(f'Invalid optimizer "{args.optimizer}"')
+
+    return optimizer