Increase OpusTrainer batch size (#1243)

evgenyrp · web-flow · commit dee8546cbb15 · 2025-09-15T14:47:14.000-07:00
diff --git a/pipeline/train/train.py b/pipeline/train/train.py
@@ -27,6 +27,7 @@
 
 
 CJK_LANGS = ["zh", "ja", "ko"]
+OPUS_TRAINER_CHUNK_SIZE = 128
 
 
 class ModelType(Enum):
@@ -390,6 +391,8 @@ def get_opustrainer_cmd(self):
                     "config": self.opustrainer_config,
                     "log-file": self.artifacts / "opustrainer.log",
                     "log-level": "INFO",
+                    "batch-size": str(OPUS_TRAINER_CHUNK_SIZE * os.cpu_count()),
+                    "chunk-size": str(OPUS_TRAINER_CHUNK_SIZE),
                 }
             ),
         ]

Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,7 @@`
`27`	`27`
`28`	`28`
`29`	`29`	`CJK_LANGS = ["zh", "ja", "ko"]`
	`30`	`+OPUS_TRAINER_CHUNK_SIZE = 128`
`30`	`31`
`31`	`32`
`32`	`33`	`class ModelType(Enum):`
`@@ -390,6 +391,8 @@ def get_opustrainer_cmd(self):`
`390`	`391`	`"config": self.opustrainer_config,`
`391`	`392`	`"log-file": self.artifacts / "opustrainer.log",`
`392`	`393`	`"log-level": "INFO",`
	`394`	`+ "batch-size": str(OPUS_TRAINER_CHUNK_SIZE * os.cpu_count()),`
	`395`	`+ "chunk-size": str(OPUS_TRAINER_CHUNK_SIZE),`
`393`	`396`	`}`
`394`	`397`	`),`
`395`	`398`	`]`