完整部署快照

ixformer_sdk/train/speedformer/speedformer.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+from abc import ABC
+from ixformer.train.speedformer.model_replacer_mapping import ModelMapping
+
+# 外部接口
+class SpeedFormer(ABC):
+    def __init__(self) -> None:
+        super().__init__()
+        self.replacer = None
+
+
+    def accelerate(self, model):
+        if model.config.model_type in ModelMapping:
+            self.replacer = ModelMapping[model.config.model_type]()
+            accelerate_model = self.replacer.accelerate(model)
+        else:
+            Warning(f"Warning: model '{model.config.model_type}' is not supported now.")
+            accelerate_model = model
+
+        return accelerate_model
+    
+    
+    def post_process(self, model):
+        self.replacer.post_process(model)