Search code examples
pythonpython-poetry

How to toggle torch cpu and gpu installation using poetry


I am trying to enable the installation of cpu and gpu versions of torch and torchvision, using poetry install --with cpu and poetry install --with gpu, respectively. I have the following in my pyproject.toml:

[tool.poetry.dependencies]
python = "^3.8"
filterpy = "^1.4.5"
gdown = "^5.1.0"
lapx = "^0.5.5"
loguru = "^0.7.2"
numpy = "1.24.4"
pyyaml = "^6.0.1"
regex = "^2024.0.0"
yacs = "^0.1.8"
scikit-learn = "^1.3.0"
pandas = "^2.0.0"
opencv-python = "^4.7.0"
ftfy = "^6.1.3"
gitpython = "^3.1.42"

[tool.poetry.group.gpu]
optional = true

[tool.poetry.group.gpu.dependencies]
torch = [
    {version = "^2.2.1", source="pypi", markers = "sys_platform == 'darwin'"},
    {version = "^2.2.1", source="torch_cuda121", markers = "sys_platform == 'linux'"},
    {version = "^2.2.1", source="torch_cuda121", markers = "sys_platform == 'win32'"},
]
torchvision = [
    {version = "^0.17.1", source="pypi", markers = "sys_platform == 'darwin'"},
    {version = "^0.17.1", source="torch_cuda121", markers = "sys_platform == 'linux'"},
    {version = "^0.17.1", source="torch_cuda121", markers = "sys_platform == 'win32'"},
]

[tool.poetry.group.cpu]
optional = true

[tool.poetry.group.cpu.dependencies]
torch = [
    {version = "^2.2.1", source="pypi", markers = "sys_platform == 'darwin'"},
    {version = "^2.2.1", source="torchcpu", markers = "sys_platform == 'linux'"},
    {version = "^2.2.1", source="torchcpu", markers = "sys_platform == 'win32'"},
]
torchvision = [
    {version = "^0.17.1", source="pypi", markers = "sys_platform == 'darwin'"},
    {version = "^0.17.1", source="torchcpu", markers = "sys_platform == 'linux'"},
    {version = "^0.17.1", source="torchcpu", markers = "sys_platform == 'win32'"},
]

[[tool.poetry.source]]
name = "torchcpu"
url = "https://download.pytorch.org/whl/cpu"
priority = "explicit"

[[tool.poetry.source]]
name = "torch_cuda121"
url = "https://download.pytorch.org/whl/cu121"
priority = "explicit"

When running poetry lock I get:

Incompatible constraints in requirements of boxmot (10.0.71):
torch (>=2.2.1,<3.0.0) ; sys_platform == "linux" or sys_platform == "win32" ; source=torchcpu
torch (>=2.2.1,<3.0.0) ; sys_platform == "linux" or sys_platform == "win32" ; source=torch_cuda121

So it seems that my way of approaching this problem does not seem to be the right one, but I have no idea how to proceed either. What am I missing? How can this be achieved with poetry?


Solution

  • The GitHub issue #6409 underscores the challenge of managing CPU and GPU dependencies in Poetry due to its lack of support for mutually exclusive groups. The proposed workarounds include using environment markers, maintaining separate configuration files, or automating the process with scripts. These methods, while not perfect, provide viable solutions to toggle between CPU and GPU installations of PyTorch and torchvision. I will provide a summary of those:

    1. Environment Markers:
    One way to manage CPU and GPU installations is to use environment markers. This method involves setting custom environment variables to indicate which version of the libraries to install. Here is an example pyproject.toml:

    [tool.poetry.dependencies]
    python = "^3.8"
    # Other dependencies
    torch = [
        {version = "^2.2.1", source="pypi", markers = "sys_platform == 'darwin' and extra == 'gpu'"},
        {version = "^2.2.1", source="torch_cuda121", markers = "sys_platform == 'linux' and extra == 'gpu'"},
        {version = "^2.2.1", source="torch_cuda121", markers = "sys_platform == 'win32' and extra == 'gpu'"},
        {version = "^2.2.1", source="pypi", markers = "sys_platform == 'darwin' and extra == 'cpu'"},
        {version = "^2.2.1", source="torchcpu", markers = "sys_platform == 'linux' and extra == 'cpu'"},
        {version = "^2.2.1", source="torchcpu", markers = "sys_platform == 'win32' and extra == 'cpu'"},
    ]
    torchvision = [
        {version = "^0.17.1", source="pypi", markers = "sys_platform == 'darwin' and extra == 'gpu'"},
        {version = "^0.17.1", source="torch_cuda121", markers = "sys_platform == 'linux' and extra == 'gpu'"},
        {version = "^0.17.1", source="torch_cuda121", markers = "sys_platform == 'win32' and extra == 'gpu'"},
        {version = "^0.17.1", source="pypi", markers = "sys_platform == 'darwin' and extra == 'cpu'"},
        {version = "^0.17.1", source="torchcpu", markers = "sys_platform == 'linux' and extra == 'cpu'"},
        {version = "^0.17.1", source="torchcpu", markers = "sys_platform == 'win32' and extra == 'cpu'"},
    ]
    
    [tool.poetry.extras]
    cpu = ["torch", "torchvision"]
    gpu = ["torch", "torchvision"]
    
    [[tool.poetry.source]]
    name = "torchcpu"
    url = "https://download.pytorch.org/whl/cpu"
    priority = "explicit"
    
    [[tool.poetry.source]]
    name = "torch_cuda121"
    url = "https://download.pytorch.org/whl/cu121"
    priority = "explicit"
    

    2. Separate pyproject.toml Files:

    Another approach is to maintain separate pyproject.toml files for CPU and GPU configurations. This method requires manual switching but ensures that only the necessary dependencies are installed.

    CPU Configuration (pyproject_cpu.toml):

    [tool.poetry.dependencies]
    python = "^3.8"
    torch = {version = "^2.2.1", source="torchcpu"}
    torchvision = {version = "^0.17.1", source="torchcpu"}
    
    [[tool.poetry.source]]
    name = "torchcpu"
    url = "https://download.pytorch.org/whl/cpu"
    priority = "explicit"
    

    GPU Configuration (pyproject_gpu.toml):

    [tool.poetry.dependencies]
    python = "^3.8"
    torch = {version = "^2.2.1", source="torch_cuda121"}
    torchvision = {version = "^0.17.1", source="torch_cuda121"}
    
    [[tool.poetry.source]]
    name = "torch_cuda121"
    url = "https://download.pytorch.org/whl/cu121"
    priority = "explicit"
    
    

    You can switch between these configurations by renaming the appropriate file to pyproject.toml before running poetry install.

    3. Custom Script: A more automated solution involves using a script to toggle between CPU and GPU configurations. Example Python Script:

    import os
    import sys
    
    def switch_dependencies(version):
        if version == 'cpu':
            os.rename('pyproject_cpu.toml', 'pyproject.toml')
        elif version == 'gpu':
            os.rename('pyproject_gpu.toml', 'pyproject.toml')
        else:
            print("Invalid version. Choose 'cpu' or 'gpu'.")
            sys.exit(1)
    
        os.system('poetry install')
    
    if __name__ == "__main__":
        if len(sys.argv) != 2:
            print("Usage: python switch_dependencies.py [cpu|gpu]")
            sys.exit(1)
        
        version = sys.argv[1]
        switch_dependencies(version)
    

    Run the script with:

    python switch_dependencies.py cpu
    

    or

    python switch_dependencies.py gpu