1

我正在运行一个实验管道来使用 PyTorch 和 CUDA 训练我的模型。我创建的环境如下：

    env = Environment.from_conda_specification(model, join(model, 'conda_dependencies.yml'))
    env.docker.enabled = True
    env.environment_variables = {'MODEL_NAME': model, 'BRANCH': branch, 'COMMIT': commit}
    env.docker.base_image = DEFAULT_GPU_IMAGE

    run_config = RunConfiguration()
    run_config.environment = env
    run_config.docker = DockerConfiguration(use_docker=True)

这是训练步骤：

train_step = PythonScriptStep(
      name='Model Train',
      source_directory=training_dir,
      compute_target=cluster,
      runconfig=run_config,
      script_name='train_aml.py',
      arguments=[
        '--model', model,
        '--model_output_dir', model_output_dir,
      ],
      inputs=[train_dataset.as_mount()],
      outputs=[model, model_output_dir]
    )

即使我在Standard_NC12_Promo运行训练脚本时使用的是机器，PyTorch 也不会拾取 GPU device = torch.device("cuda" if torch.cuda.is_available() else "cpu")。
如果我尝试在同一台机器上而不是在实验中运行我的脚本，则使用 GPU。
你知道任何潜在的解决方案吗？

4

1 回答 1

1

根据 pytorch 版本，您可能需要特定版本的 cuda。从这里尝试 cuda 11.0.3 或 cuda 11.1 https://github.com/Azure/AzureML-Containers/tree/master/base/gpu

关于您的代码片段，请将环境变量移出环境对象以运行配置

于 2021-08-09T13:15:09.217 回答