adds TPU to CI.

.github/workflows/tpu-tests.yaml

@@ -17,3 +17,43 @@ jobs:
         conda info -a
         conda create -y -n py36 python=3.6
         source activate py36
+
+
+    - name: Install Torch XLA and others
+      run: |
+        export PATH="$HOME/miniconda/bin:$PATH"
+        source activate py36
+
+        ## Install gsutil
+        echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | sudo tee -a /etc/apt/sources.list.d/google-cloud-sdk.list
+        sudo apt-get install -y apt-transport-https ca-certificates gnupg curl
+        curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -
+        sudo apt-get update && sudo apt-get install -y google-cloud-sdk
+
+        ## Install openblas and mkl
+        sudo apt-get install -y libopenblas-dev
+        conda install -y mkl=2020.0=166  # if change the version, think about to change LD_LIBRARY_PATH below
+
+        ## Download torch & xla
+        gsutil cp gs://tpu-pytorch/wheels/torch-1.5-cp36-cp36m-linux_x86_64.whl .
+        gsutil cp gs://tpu-pytorch/wheels/torch_xla-1.5-cp36-cp36m-linux_x86_64.whl .
+
+        ## Install torch & xla
+        pip install torch-1.5-cp36-cp36m-linux_x86_64.whl
+        pip install torch_xla-1.5-cp36-cp36m-linux_x86_64.whl
+
+        ## Install test deps and Ignite
+        pip install -r requirements-dev.txt
+        python setup.py install
+
+    - name: Run Tests
+      run: |
+        export PATH="$HOME/miniconda/bin:$PATH"
+        source activate py36
+
+        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/miniconda/pkgs/mkl-2020.0-166/lib/
+        export XRT_DEVICE_MAP="CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0"
+        export XRT_WORKERS="localservice:0;grpc://localhost:40934"
+
+        python -c "import torch_xla; print('torch xla version:', torch_xla.__version__)"
+        py.test --cov ignite --cov-append --cov-report term-missing tests/ -vvv -m tpu

setup.cfg

@@ -16,4 +16,5 @@ ignore = E203,E231,E305,E402,E721,E722,E741,F401,F403,F405,F821,F841,F999,W503
 [tool:pytest]
 markers =
     distributed: mark a test with distributed option
-    multinode_distributed: mark a test with multi-node distributed option
+    multinode_distributed: mark a test with multi-node distributed option
+    tpu: mark a test as requiring XLA

tests/ignite/engine/test_create_supervised.py

@@ -10,6 +10,11 @@
 from ignite.engine.engine import Events
 from ignite.metrics import MeanSquaredError
 
+try:
+    import torch_xla.core.xla_model as xm
+    has_xla = True
+except ImportError:
+    has_xla = False
 
 def test_create_supervised_trainer():
     model = Linear(1, 1)
@@ -103,6 +108,30 @@ def test_create_supervised_trainer_on_cuda():
     assert model.bias.item() == approx(0.8)
 
 
+@pytest.mark.skipif(not has_xla, reason="Skip if no TPU")
+def test_create_supervised_trainer_on_tpu():
+    device = "xla"
+    model = Linear(1, 1)
+    model.to(device)
+    model.weight.data.zero_()
+    model.bias.data.zero_()
+    optimizer = SGD(model.parameters(), 0.1)
+    trainer = create_supervised_trainer(model, optimizer, mse_loss, device=device)
+
+    x = torch.tensor([[1.0], [2.0]])
+    y = torch.tensor([[3.0], [5.0]])
+    data = [(x, y)]
+
+    assert model.weight.data[0, 0].item() == approx(0.0)
+    assert model.bias.item() == approx(0.0)
+
+    state = trainer.run(data)
+
+    assert state.output == approx(17.0)
+    assert model.weight.data[0, 0].item() == approx(1.3)
+    assert model.bias.item() == approx(0.8)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
 def test_create_supervised_trainer_on_cuda_with_model_on_cpu():
     model = Linear(1, 1)