Set workers=0 for MPS Train and Val modes (#4697)

This commit is contained in:
Glenn Jocher 2023-09-02 04:25:03 +02:00 committed by GitHub
parent 2bc6e647c7
commit a4fabfdacf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 88 additions and 74 deletions

View file

@ -196,17 +196,27 @@ def add_integration_callbacks(instance):
instance (Trainer, Predictor, Validator, Exporter): An object with a 'callbacks' attribute that is a dictionary
of callback lists.
"""
from .clearml import callbacks as clearml_cb
from .comet import callbacks as comet_cb
from .dvc import callbacks as dvc_cb
from .hub import callbacks as hub_cb
from .mlflow import callbacks as mlflow_cb
from .neptune import callbacks as neptune_cb
from .raytune import callbacks as tune_cb
from .tensorboard import callbacks as tensorboard_cb
from .wb import callbacks as wb_cb
for x in clearml_cb, comet_cb, hub_cb, mlflow_cb, neptune_cb, tune_cb, tensorboard_cb, wb_cb, dvc_cb:
for k, v in x.items():
if v not in instance.callbacks[k]: # prevent duplicate callbacks addition
instance.callbacks[k].append(v) # callback[name].append(func)
# Load HUB callbacks
from .hub import callbacks
# Load training callbacks
if 'Trainer' in instance.__class__.__name__:
from .clearml import callbacks as clear_cb
from .comet import callbacks as comet_cb
from .dvc import callbacks as dvc_cb
from .mlflow import callbacks as mlflow_cb
from .neptune import callbacks as neptune_cb
from .raytune import callbacks as tune_cb
from .tensorboard import callbacks as tb_cb
from .wb import callbacks as wb_cb
callbacks.update({**clear_cb, **comet_cb, **dvc_cb, **mlflow_cb, **neptune_cb, **tune_cb, **tb_cb, **wb_cb})
# Load export callbacks (patch to avoid CoreML protobuf error)
if 'Exporter' in instance.__class__.__name__:
from .tensorboard import callbacks as tb_cb
callbacks.update(tb_cb)
for k, v in callbacks.items():
if v not in instance.callbacks[k]: # prevent duplicate callbacks addition
instance.callbacks[k].append(v) # callback[name].append(func)