首页 > 解决方案 > python setup.py:安装 tensorflow 还是 tensorflow-gpu?

问题描述

我正在设计一个需要tensorflow. 由于 TensorFlow 有多个安装(tensorflowtensorflow-gpu),我如何在我的install_requires部分中添加任何一个都可以?

标签: pythonpython-3.xsetuptoolssetup.py

解决方案


所以我是这样解决的:

from pkg_resources import DistributionNotFound, get_distribution
from setuptools import setup, find_packages    

def get_dist(pkgname):
    try:
        return get_distribution(pkgname)
    except DistributionNotFound:
        return None

install_deps = ['numpy', 'tensorflow']
if get_dist('tensorflow') is None and get_dist('tensorflow-gpu') is not None:
    install_deps.remove('tensorflow')

setup(..., install_requires=install_deps)

推荐阅读