support cpu training, use cpu training on mac

This commit is contained in:
Lion
2024-03-13 22:09:24 +08:00
parent 931781774d
commit 1963eb01cc
8 changed files with 41 additions and 43 deletions

View File

@@ -118,16 +118,16 @@ def main(args):
os.environ["MASTER_ADDR"]="localhost"
trainer: Trainer = Trainer(
max_epochs=config["train"]["epochs"],
accelerator="gpu",
accelerator="gpu" if torch.cuda.is_available() else "cpu",
# val_check_interval=9999999999999999999999,###不要验证
# check_val_every_n_epoch=None,
limit_val_batches=0,
devices=-1,
devices=-1 if torch.cuda.is_available() else 1,
benchmark=False,
fast_dev_run=False,
strategy = "auto" if torch.backends.mps.is_available() else DDPStrategy(
strategy = DDPStrategy(
process_group_backend="nccl" if platform.system() != "Windows" else "gloo"
), # mps 不支持多节点训练
) if torch.cuda.is_available() else "auto",
precision=config["train"]["precision"],
logger=logger,
num_sanity_val_steps=0,