Como verificar se meu código python consegue usar a GPU Nvidia

Com esse código simples em Python, conseguimos identificar se o mesmo está apto a usar a GPU para acelerar cálculos intensivos, como treinamento de redes neurais ou processamento de imagens, etc..

Antes de executar o código, tenha certeza de que os driver de sua gpu estão instalados e configurados corretamente.

Instale os módulos python necessários:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

Crie um arquivo chamado `test-gpu.py` com o seguinte conteúdo:

import torch
import subprocess
import platform

def list_gpus():
    print("🔍 Verificando GPUs disponíveis...\n")
    
    # 1. Verifica se PyTorch está disponível (mais comum em ambientes Jupyter com GPU)
    if torch.cuda.is_available():
        print("GPU(s) disponíveis (via PyTorch):")
        for i in range(torch.cuda.device_count()):
            device = torch.cuda.get_device_properties(i)
            print(f"GPU {i}: {device.name} | Memória: {device.total_memory / 1024**2:.2f} MB")
        return

    # 2. Verifica se `nvidia-smi` está disponível (sistema Linux com drivers NVIDIA)
    if platform.system() == "Linux":
        try:
            result = subprocess.run(["nvidia-smi", "-q", "-d", "GPU"], capture_output=True, text=True, check=True)
            print("GPU(s) disponíveis (via nvidia-smi):")
            print(result.stdout)
        except subprocess.CalledProcessError as e:
            print("⚠️ Erro ao executar nvidia-smi:", e.stderr)
        except FileNotFoundError:
            print("⚠️ Comando `nvidia-smi` não encontrado. Verifique se os drivers NVIDIA estão instalados.")
        return

    # 3. Caso nenhuma GPU seja detectada
    print("⚠️ Nenhuma GPU detectada. Verifique se o sistema tem suporte a GPU e as ferramentas adequadas instaladas.")

if __name__ == "__main__":
    list_gpus()

Resultado esperado:

🔍 Verificando GPUs disponíveis...

GPU(s) disponíveis (via PyTorch):
GPU 0: NVIDIA GeForce RTX 3080 | Memória: 10002.44 MB