python - 为什么 Databricks 笔记本单元一直在训练我的 ML 模型“正在运行命令……”？

Question

我正在使用 PyTorch Lightning 在 Azure Databricks 集群（在笔记本上）上训练 PyTorch 模型，并使用 mlflow 进行跟踪。

我想在 Databricks 托管的跟踪服务器上存储训练指标 + 工件。

要启用它，代码如下：

mlflow.pytorch.autolog()
trainer = pl.Trainer(gpus=1, max_epochs=30, callbacks=[EarlyStopping(monitor='val_loss', patience = 6)], progress_bar_refresh_rate=0)
trainer.fit(classifier, train_dl, valid_dl)
print("Done")

但是，笔记本单元格在“运行命令...”状态中停留的时间太长：

即使在驱动程序日志中执行似乎已经结束：

并且实验在 mlflow UI 中标记为 FINISHED：

手动停止执行也无法解决，因为单元格将永远处于“正在取消...”状态。所以剩下的唯一选择就是清除集群状态。

这是一个问题，因为我无法执行对工件日志记录有用的进一步命令：

mapping.to_json("/tmp/mapping.json", orient = "records")
mlflow.log_artifact("/tmp/mapping.json", "mapping")

torch.save(classifier.state_dict(), "/tmp/model.pt")
mlflow.log_artifact("/tmp/model.pt", "model.pt")

这个问题似乎与 GC 问题有关：

2021-04-22T07:58:58.025+0000: [GC (分配失败) [PSYoungGen: 28399104K->100290K(28499456K)] 28696488K->397698K(85935104K), 0.0755720 secs] [Times:06, user=0.真实=0.08 秒] 2021-04-22T08:01:01.645+0000: [GC (System.gc()) [PSYoungGen: 4522724K->54360K(28561920K)] 4820132K->351776K(85997568K), 0.0237712 秒] ：用户=0.09 sys=0.01，真实=0.02 秒] 2021-04-22T08:01:01.669+0000：[完整 GC (System.gc()) [PSYoungGen: 54360K->0K(28561920K)] [ParOldGen: 297416K ->123173K(57435648K)] 351776K->123173K(85997568K), [元空间: 203356K->203325K(219136K)], 0.3513905 秒] [时间: user=0.99 sys=0.00, real=0.36 secs]

难道我做错了什么？我应该以其他方式跟踪实验 + 工件吗？运行单节点和集群都不起作用，也不能减少训练集的大小。

标签： pythonpytorchazure-databricksmlflowpytorch-lightning