pytorch 如何在Vertex AI Pipeline中启用GPU

lx0bsm1f 于 6个月前发布在其他

关注(0)|答案(1)|浏览(75)

我是Vertex AI Pipeline的新手，我遇到了一个问题，我无法在我的管道内启用GPU利用率。尽管为特定组件指定了硬件类型，但它仍然在CPU上运行，如下图所示。我怀疑我的管道设置中可能缺少配置。

embeddings_task = (
        generate_article_embeddings(
            transformer_model=TRANSFORMER_MODEL,
            article_dataset=articles.output,
        )
        .set_cpu_limit("32")
        .set_memory_limit("208G")
        .add_node_selector_constraint("NVIDIA_TESLA_T4")
        .set_gpu_limit("4")
    )

字符串
enter image description here
有没有人可以提供关于如何在Vertex AI Pipeline中正确启用GPU利用率的指导？
如有任何帮助，我们将不胜感激。
我已经安装了所有必要的GPU库来启用GPU加速，但当我运行代码时，它仍然默认使用CPU。

pytorch

来源：https://stackoverflow.com/questions/77034770/how-to-enable-gpu-in-vertex-ai-pipeline

1条答案

按热度按时间

bgtovc5b1#

你需要确保你已经在你的基本容器中安装了cuda依赖。
我制作这个组件是为了测试GPU是否可以使用Tensor：

@dsl.component
def simple(
    word: str,
    number: int,
) -> None:
    log.info("Running Simple")

    if torch.cuda.is_available():
        device = torch.device("cuda")  # set the device to be the first cuda device
        log.info("PyTorch is using the GPU")
        log.info("Current device: %s", torch.cuda.current_device())
        log.info("Device name: %s", torch.cuda.get_device_name(device))
    else:
        log.info("PyTorch is using the CPU")

    # Create a tensor on the CPU
    x = torch.randn(100, number)

    if torch.cuda.is_available():
        # Move the tensor to the GPU
        x = x.to(device)

    # Perform some computations on the tensor on the GPU
    y = x * 10**4

    # Move the tensor back to the CPU for printing
    y = y.to("cpu")

字符串

赞(0）回复(0）举报 6个月前

我来回答

pytorch 如何在Vertex AI Pipeline中启用GPU

1条答案

相关问题

热门标签

最新问答