Skip to content

Instantly share code, notes, and snippets.

@HeinrichAD
Last active November 27, 2020 21:46
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save HeinrichAD/eea00d50c5702cbf63082983b6a2d763 to your computer and use it in GitHub Desktop.
Save HeinrichAD/eea00d50c5702cbf63082983b6a2d763 to your computer and use it in GitHub Desktop.
Google Colab Processor Unit / Runtime Type Detector.
Display the source blob
Display the rendered blob
Raw
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "colab_pu_detector.ipynb",
"provenance": [],
"collapsed_sections": []
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "TPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "gwLFwcWubtA4"
},
"source": [
"# Google Colab Processor Unit / Runtime Type\n",
"\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/gist/HeinrichAD/eea00d50c5702cbf63082983b6a2d763/colab_pu_detector.ipynb)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hoxAOOacmklJ"
},
"source": [
"Import via `httpimport` directly from GitHub or copy the mini [`ProcessorUnit`](https://gist.github.com/HeinrichAD/eea00d50c5702cbf63082983b6a2d763#file-processorunit-py) class into your notebook."
]
},
{
"cell_type": "code",
"metadata": {
"id": "3sXDPH9tFaA2"
},
"source": [
"import httpimport\n",
"\n",
"url = \"https://gist.githubusercontent.com/HeinrichAD/eea00d50c5702cbf63082983b6a2d763/raw/48a281f4d5738928e73f53f59364ea2d622f050e\"\n",
"with httpimport.remote_repo([\"ProcessorUnit\"], url):\n",
" from ProcessorUnit import PU"
],
"execution_count": 1,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "K4tBDgqPa7gh"
},
"source": [
"## Examples\n",
"\n",
"Get current Processor Unit"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "WjvyK6AWhQGg",
"outputId": "97dfd604-5498-41a8-9925-a625119dfb85"
},
"source": [
"print(PU.getPU())"
],
"execution_count": 2,
"outputs": [
{
"output_type": "stream",
"text": [
"PU.TPU\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xEcBJUrOhU5Y"
},
"source": [
"Check specific Processor Unit"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "0jfC9vMNhXoE",
"outputId": "3ffad5c9-800e-4e6f-86f8-da431f36f58d"
},
"source": [
"print(PU.isTPU())"
],
"execution_count": 3,
"outputs": [
{
"output_type": "stream",
"text": [
"True\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "lobIvny0hall"
},
"source": [
"The usual Python Enum power"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "ZBrf6zlWhZ4T",
"outputId": "166d92d2-12f0-44b5-94b3-73723a4a7163"
},
"source": [
"for pu in PU:\n",
" print(pu)"
],
"execution_count": 4,
"outputs": [
{
"output_type": "stream",
"text": [
"PU.CPU\n",
"PU.GPU\n",
"PU.TPU\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RipMpUOpc8Y6"
},
"source": [
"## Basis of the Processor Unit decision\n",
"\n",
"```python\n",
"from os import environ\n",
"print(environ[\"COLAB_GPU\"] == '1', \"TPU_NAME\" in environ)\n",
"```\n",
"```\n",
"If ...\n",
" False False CPU\n",
" True False GPU\n",
" False True TPU\n",
"```"
]
}
]
}
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment