Google Colab Processor Unit / Runtime Type Detector.

colab_pu_detector.ipynb

ProcessorUnit.py

from enum import Enum
from os import environ

class PU(Enum):
  """
  Google Colab Processor Unit / Runtime Type Detector.
  """
  CPU = 1
  GPU = 2
  TPU = 3

  def isCPU():
    return not (PU.isGPU() or PU.isTPU())
  def isGPU():
    return environ["COLAB_GPU"] == '1'
  def isTPU():
    return "TPU_NAME" in environ

  def getPU():
    return PU.GPU if PU.isGPU() else PU.TPU if PU.isTPU() else PU.CPU