Find a CUDA compute capability from a GPU name programmatically

get_ccompute.py

#!/usr/bin/env python

"""
Obtain the correspondence of a local GPU name to its cuda compute capability
Names match the output of nvidia-smi --query-gpu=gpu_name --format=csv,noheader

Usage:
----
./get_ccompute.py <GPU_NAME>
"""
import sys
try:
    from urllib.request import urlopen
except ImportError:
    from urllib2 import urlopen
import collections

from bs4 import BeautifulSoup

def process(url):
    page = urlopen(url)

    text = page.read()
    page.close()

    soup = BeautifulSoup(text, "lxml")
    d = collections.defaultdict(str)

    for table in soup.findAll('table', {'class': 'table table-striped'}):
        for line in table.findAll('tr'):
            gpu = line.findNext('td').text.strip()
            compute = line.findNext('td').findNext('td').text.strip()
            d[gpu] = compute
    return d

def main():
    if len(sys.argv) > 1:
        gpu = sys.argv[1]
    if len(sys.argv) == 1:
        print("Usage: %s <GPU_NAME> ..." % sys.argv[0])
        sys.exit(1)
    # at least one gpu was passed
    d = process("https://developer.nvidia.com/cuda-gpus")
    print(d[gpu])

if __name__ == "__main__":
    main()

Tweaked to use python3, and added more useful Usage. Only 2 non-comment lines changed: #!/usr/bin/env python3 and

        print("Usage: %s "   ' "$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader)" '   % sys.argv[0])

Would be nice to add URL to this gist in comments at top, too...

"""
Obtain the correspondence of a local <GPU_NAME> to its cuda compute capability
<GPU_NAME>  can be output of:  nvidia-smi --query-gpu=gpu_name --format=csv,noheader
https://gist.github.com/huitseeker/b2c79e5b763d58b06b9985de2b3c0d4d

Usage:
----
./get_ccompute.py <GPU_NAME> 
Example: ./get_ccompute.py "GeForce GTX 965M"
Note: <GPU NAME> is output of: nvidia-smi --query-gpu=gpu_name --format=csv,noheader
./get_ccompute.py "$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader)" # should work
"""