demo_application.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.0"
    },
    "colab": {
      "name": "demo application.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/walexi/pytorch_challenge/blob/master/demo_application.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QyPkwvwIx_IM",
        "colab_type": "text"
      },
      "source": [
        "## Encrypted ML in Health Care\n",
        "In this quick demo, we show the distributed training of a the Graph Attention Networks (GATConv from Veličković et al.: [Graph Attention Networks](https://arxiv.org/abs/1710.10903) (ICLR 2018)) on a dataset (to simulate private health data(1))\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "\n",
        "*1 The protein-protein interaction networks from the `\"Predicting\n",
        "    Multicellular Function through Multi-layer Tissue Networks\"\n",
        "    <https://arxiv.org/abs/1707.04638>`_ paper, containing positional gene\n",
        "    sets, motif gene sets and immunological signatures as features (50 in\n",
        "    total) and gene ontology sets as labels (121 in total).*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yMfy9kGrmtV6",
        "colab_type": "text"
      },
      "source": [
        "## Imports and training configuration"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "epBgLzMfmtV8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torchvision import datasets, transforms\n",
        "import time\n",
        "import os.path as osp\n",
        "from sklearn.metrics import f1_score\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tQtVru5Jnaoo",
        "colab_type": "code",
        "outputId": "bc83c606-4946-4700-9329-80a700bca38c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "!pip install --verbose --no-cache-dir torch-scatter\n",
        "!pip install --verbose --no-cache-dir torch-sparse\n",
        "!pip install --verbose --no-cache-dir torch-cluster\n",
        "!pip install --verbose --no-cache-dir torch-spline-conv\n",
        "!pip install torch-geometric"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cpu/graclus.o -o build/lib.linux-x86_64-3.6/torch_cluster/graclus_cpu.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.grid_cpu' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/include/python3.6m -c cpu/grid.cpp -o build/temp.linux-x86_64-3.6/cpu/grid.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=grid_cpu -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cpu/grid.o -o build/lib.linux-x86_64-3.6/torch_cluster/grid_cpu.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.fps_cpu' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/include/python3.6m -c cpu/fps.cpp -o build/temp.linux-x86_64-3.6/cpu/fps.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=fps_cpu -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cpu/fps.o -o build/lib.linux-x86_64-3.6/torch_cluster/fps_cpu.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.rw_cpu' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/include/python3.6m -c cpu/rw.cpp -o build/temp.linux-x86_64-3.6/cpu/rw.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=rw_cpu -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cpu/rw.o -o build/lib.linux-x86_64-3.6/torch_cluster/rw_cpu.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.sampler_cpu' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/include/python3.6m -c cpu/sampler.cpp -o build/temp.linux-x86_64-3.6/cpu/sampler.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=sampler_cpu -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cpu/sampler.o -o build/lib.linux-x86_64-3.6/torch_cluster/sampler_cpu.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.graclus_cuda' extension\n",
            "  creating build/temp.linux-x86_64-3.6/cuda\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/graclus.cpp -o build/temp.linux-x86_64-3.6/cuda/graclus.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=graclus_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/graclus_kernel.cu -o build/temp.linux-x86_64-3.6/cuda/graclus_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=graclus_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cuda/graclus.o build/temp.linux-x86_64-3.6/cuda/graclus_kernel.o -L/usr/local/cuda/lib64 -lcudart -o build/lib.linux-x86_64-3.6/torch_cluster/graclus_cuda.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.grid_cuda' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/grid.cpp -o build/temp.linux-x86_64-3.6/cuda/grid.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=grid_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/grid_kernel.cu -o build/temp.linux-x86_64-3.6/cuda/grid_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=grid_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cuda/grid.o build/temp.linux-x86_64-3.6/cuda/grid_kernel.o -L/usr/local/cuda/lib64 -lcudart -o build/lib.linux-x86_64-3.6/torch_cluster/grid_cuda.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.fps_cuda' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/fps.cpp -o build/temp.linux-x86_64-3.6/cuda/fps.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=fps_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/fps_kernel.cu -o build/temp.linux-x86_64-3.6/cuda/fps_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=fps_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cuda/fps.o build/temp.linux-x86_64-3.6/cuda/fps_kernel.o -L/usr/local/cuda/lib64 -lcudart -o build/lib.linux-x86_64-3.6/torch_cluster/fps_cuda.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.nearest_cuda' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/nearest.cpp -o build/temp.linux-x86_64-3.6/cuda/nearest.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=nearest_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/nearest_kernel.cu -o build/temp.linux-x86_64-3.6/cuda/nearest_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=nearest_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cuda/nearest.o build/temp.linux-x86_64-3.6/cuda/nearest_kernel.o -L/usr/local/cuda/lib64 -lcudart -o build/lib.linux-x86_64-3.6/torch_cluster/nearest_cuda.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.knn_cuda' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/knn.cpp -o build/temp.linux-x86_64-3.6/cuda/knn.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=knn_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/knn_kernel.cu -o build/temp.linux-x86_64-3.6/cuda/knn_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=knn_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cuda/knn.o build/temp.linux-x86_64-3.6/cuda/knn_kernel.o -L/usr/local/cuda/lib64 -lcudart -o build/lib.linux-x86_64-3.6/torch_cluster/knn_cuda.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.radius_cuda' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/radius.cpp -o build/temp.linux-x86_64-3.6/cuda/radius.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=radius_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/radius_kernel.cu -o build/temp.linux-x86_64-3.6/cuda/radius_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=radius_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cuda/radius.o build/temp.linux-x86_64-3.6/cuda/radius_kernel.o -L/usr/local/cuda/lib64 -lcudart -o build/lib.linux-x86_64-3.6/torch_cluster/radius_cuda.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_cluster.rw_cuda' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/rw.cpp -o build/temp.linux-x86_64-3.6/cuda/rw.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=rw_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/rw_kernel.cu -o build/temp.linux-x86_64-3.6/cuda/rw_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=rw_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cuda/rw.o build/temp.linux-x86_64-3.6/cuda/rw_kernel.o -L/usr/local/cuda/lib64 -lcudart -o build/lib.linux-x86_64-3.6/torch_cluster/rw_cuda.cpython-36m-x86_64-linux-gnu.so\n",
            "  installing to build/bdist.linux-x86_64/wheel\n",
            "  running install\n",
            "  running install_lib\n",
            "  creating build/bdist.linux-x86_64\n",
            "  creating build/bdist.linux-x86_64/wheel\n",
            "  creating build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/utils.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/__init__.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_radius.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_sampler.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_grid.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_nearest.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_graclus.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_rw.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_knn.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_fps.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  creating build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/__init__.py -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/sampler_cpu.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/radius.py -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/fps_cpu.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/radius_cuda.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/grid.py -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/knn_cuda.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/nearest.py -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/rw.py -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/grid_cpu.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/nearest_cuda.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/graclus_cuda.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/rw_cuda.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/fps.py -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/sampler.py -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/graclus_cpu.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/grid_cuda.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/knn.py -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/graclus.py -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/rw_cpu.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  copying build/lib.linux-x86_64-3.6/torch_cluster/fps_cuda.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_cluster\n",
            "  running install_egg_info\n",
            "  running egg_info\n",
            "  writing torch_cluster.egg-info/PKG-INFO\n",
            "  writing dependency_links to torch_cluster.egg-info/dependency_links.txt\n",
            "  writing requirements to torch_cluster.egg-info/requires.txt\n",
            "  writing top-level names to torch_cluster.egg-info/top_level.txt\n",
            "  reading manifest file 'torch_cluster.egg-info/SOURCES.txt'\n",
            "  reading manifest template 'MANIFEST.in'\n",
            "  writing manifest file 'torch_cluster.egg-info/SOURCES.txt'\n",
            "  Copying torch_cluster.egg-info to build/bdist.linux-x86_64/wheel/torch_cluster-1.4.4-py3.6.egg-info\n",
            "  running install_scripts\n",
            "  adding license file \"LICENSE\" (matched pattern \"LICEN[CS]E*\")\n",
            "  creating build/bdist.linux-x86_64/wheel/torch_cluster-1.4.4.dist-info/WHEEL\n",
            "  creating '/tmp/pip-wheel-4h0uvhdd/torch_cluster-1.4.4-cp36-cp36m-linux_x86_64.whl' and adding 'build/bdist.linux-x86_64/wheel' to it\n",
            "  adding 'test/__init__.py'\n",
            "  adding 'test/test_fps.py'\n",
            "  adding 'test/test_graclus.py'\n",
            "  adding 'test/test_grid.py'\n",
            "  adding 'test/test_knn.py'\n",
            "  adding 'test/test_nearest.py'\n",
            "  adding 'test/test_radius.py'\n",
            "  adding 'test/test_rw.py'\n",
            "  adding 'test/test_sampler.py'\n",
            "  adding 'test/utils.py'\n",
            "  adding 'torch_cluster/__init__.py'\n",
            "  adding 'torch_cluster/fps.py'\n",
            "  adding 'torch_cluster/fps_cpu.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/fps_cuda.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/graclus.py'\n",
            "  adding 'torch_cluster/graclus_cpu.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/graclus_cuda.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/grid.py'\n",
            "  adding 'torch_cluster/grid_cpu.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/grid_cuda.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/knn.py'\n",
            "  adding 'torch_cluster/knn_cuda.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/nearest.py'\n",
            "  adding 'torch_cluster/nearest_cuda.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/radius.py'\n",
            "  adding 'torch_cluster/radius_cuda.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/rw.py'\n",
            "  adding 'torch_cluster/rw_cpu.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/rw_cuda.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster/sampler.py'\n",
            "  adding 'torch_cluster/sampler_cpu.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_cluster-1.4.4.dist-info/LICENSE'\n",
            "  adding 'torch_cluster-1.4.4.dist-info/METADATA'\n",
            "  adding 'torch_cluster-1.4.4.dist-info/WHEEL'\n",
            "  adding 'torch_cluster-1.4.4.dist-info/top_level.txt'\n",
            "  adding 'torch_cluster-1.4.4.dist-info/RECORD'\n",
            "  removing build/bdist.linux-x86_64/wheel\n",
            "\u001b[?25hdone\n",
            "  Created wheel for torch-cluster: filename=torch_cluster-1.4.4-cp36-cp36m-linux_x86_64.whl size=15480016 sha256=52a4f5a59b01f555781429ca405e6153d904efb64a3f19dce4074680cbd4fbd5\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-2u8m36pb/wheels/20/7b/ab/b3e266920055d1e51988f93a99ef8df62e399b234c8d50527f\n",
            "  Removing source in /tmp/pip-install-y1tizdg_/torch-cluster\n",
            "Successfully built torch-cluster\n",
            "Installing collected packages: torch-cluster\n",
            "\n",
            "Successfully installed torch-cluster-1.4.4\n",
            "Cleaning up...\n",
            "Removed build tracker '/tmp/pip-req-tracker-mdw_o37w'\n",
            "Created temporary directory: /tmp/pip-ephem-wheel-cache-mzwr_9wr\n",
            "Created temporary directory: /tmp/pip-req-tracker-8vfpxnhu\n",
            "Created requirements tracker '/tmp/pip-req-tracker-8vfpxnhu'\n",
            "Created temporary directory: /tmp/pip-install-a980v7k7\n",
            "Collecting torch-spline-conv\n",
            "  1 location(s) to search for versions of torch-spline-conv:\n",
            "  * https://pypi.org/simple/torch-spline-conv/\n",
            "  Getting page https://pypi.org/simple/torch-spline-conv/\n",
            "  Found index url https://pypi.org/simple\n",
            "  Starting new HTTPS connection (1): pypi.org:443\n",
            "  https://pypi.org:443 \"GET /simple/torch-spline-conv/ HTTP/1.1\" 200 936\n",
            "  Analyzing links from page https://pypi.org/simple/torch-spline-conv/\n",
            "    Found link https://files.pythonhosted.org/packages/43/48/02fd06eca47d38efc031b34875c2e7453faab7bde7c02b2669957d230693/torch_spline_conv-0.1.0.tar.gz#sha256=b1d7bceecffb4bc394d9f5df57b56fbff95867aeb778588a054e95a6ba5cd610 (from https://pypi.org/simple/torch-spline-conv/), version: 0.1.0\n",
            "    Found link https://files.pythonhosted.org/packages/99/a6/19a54aaaf507960814af5b6ec3d6916497d54a9345d0b64a98764465fddc/torch_spline_conv-1.0.0.tar.gz#sha256=72950ae9cc15dd00d0977ffca115af0e59c10f1dbac33dd800307fef2b5808c3 (from https://pypi.org/simple/torch-spline-conv/), version: 1.0.0\n",
            "    Found link https://files.pythonhosted.org/packages/4a/e9/973d572b69f8fef3b08d86f77312a75ab4c00ffe8d3582ae09b433b34b8a/torch_spline_conv-1.0.1.tar.gz#sha256=6be33fab8402950b2ddc1dec950d0b5d4c6f183bcf53bca2534a03eb4fd5801b (from https://pypi.org/simple/torch-spline-conv/), version: 1.0.1\n",
            "    Found link https://files.pythonhosted.org/packages/bc/ff/4e834a74fd180f87f6f9aedcfd876cca4c6fa706afa31e9541838f6958d1/torch_spline_conv-1.0.3.tar.gz#sha256=a0be3bc9847458ca92146c2da754dbd2f910be426414757256416e507240da88 (from https://pypi.org/simple/torch-spline-conv/), version: 1.0.3\n",
            "    Found link https://files.pythonhosted.org/packages/1b/d2/98328ef2395b70b9795ddbd091398e49e918f8892d1ea8dad869550b684b/torch_spline_conv-1.0.4.tar.gz#sha256=26e2bb0832d45d185f2e7b38569e64e7f70af8ca2a8b05663972f261ba76a978 (from https://pypi.org/simple/torch-spline-conv/), version: 1.0.4\n",
            "    Found link https://files.pythonhosted.org/packages/92/86/f086f96e09654106ad527f4b03ac6188c1299a9fe555fc39c663b155a05a/torch_spline_conv-1.0.5.tar.gz#sha256=fffa87c94703e4c43a46bd8ec149cd6c125586ebd1f3c04fec48d1f6b60a4f3c (from https://pypi.org/simple/torch-spline-conv/), version: 1.0.5\n",
            "    Found link https://files.pythonhosted.org/packages/9a/6d/b34721af4bb907814a41a1e0e5e426c97aee644efef6877ed112bfb3d81c/torch_spline_conv-1.0.6.tar.gz#sha256=a8bd72bac7dc078ddfbf22789c5fcdf159901494c626344d67ad203e7f8c9b1e (from https://pypi.org/simple/torch-spline-conv/), version: 1.0.6\n",
            "    Found link https://files.pythonhosted.org/packages/3c/dd/daa9d0b7b2ede913e573876ae286a58ec296678858f2814ff6d6789b234f/torch_spline_conv-1.1.0.tar.gz#sha256=e6029526205d1f7cb535389bebd81decf0649a20ea6a67688c02bd335a7f9339 (from https://pypi.org/simple/torch-spline-conv/), version: 1.1.0\n",
            "  Given no hashes to check 8 links for project 'torch-spline-conv': discarding no candidates\n",
            "  Using version 1.1.0 (newest of versions: 0.1.0, 1.0.0, 1.0.1, 1.0.3, 1.0.4, 1.0.5, 1.0.6, 1.1.0)\n",
            "  Created temporary directory: /tmp/pip-unpack-fky0ydri\n",
            "  Starting new HTTPS connection (1): files.pythonhosted.org:443\n",
            "  https://files.pythonhosted.org:443 \"GET /packages/3c/dd/daa9d0b7b2ede913e573876ae286a58ec296678858f2814ff6d6789b234f/torch_spline_conv-1.1.0.tar.gz HTTP/1.1\" 200 12931\n",
            "  Downloading https://files.pythonhosted.org/packages/3c/dd/daa9d0b7b2ede913e573876ae286a58ec296678858f2814ff6d6789b234f/torch_spline_conv-1.1.0.tar.gz\n",
            "  Downloading from URL https://files.pythonhosted.org/packages/3c/dd/daa9d0b7b2ede913e573876ae286a58ec296678858f2814ff6d6789b234f/torch_spline_conv-1.1.0.tar.gz#sha256=e6029526205d1f7cb535389bebd81decf0649a20ea6a67688c02bd335a7f9339 (from https://pypi.org/simple/torch-spline-conv/)\n",
            "  Added torch-spline-conv from https://files.pythonhosted.org/packages/3c/dd/daa9d0b7b2ede913e573876ae286a58ec296678858f2814ff6d6789b234f/torch_spline_conv-1.1.0.tar.gz#sha256=e6029526205d1f7cb535389bebd81decf0649a20ea6a67688c02bd335a7f9339 to build tracker '/tmp/pip-req-tracker-8vfpxnhu'\n",
            "    Running setup.py (path:/tmp/pip-install-a980v7k7/torch-spline-conv/setup.py) egg_info for package torch-spline-conv\n",
            "    Running command python setup.py egg_info\n",
            "    running egg_info\n",
            "    creating pip-egg-info/torch_spline_conv.egg-info\n",
            "    writing pip-egg-info/torch_spline_conv.egg-info/PKG-INFO\n",
            "    writing dependency_links to pip-egg-info/torch_spline_conv.egg-info/dependency_links.txt\n",
            "    writing top-level names to pip-egg-info/torch_spline_conv.egg-info/top_level.txt\n",
            "    writing manifest file 'pip-egg-info/torch_spline_conv.egg-info/SOURCES.txt'\n",
            "    reading manifest file 'pip-egg-info/torch_spline_conv.egg-info/SOURCES.txt'\n",
            "    reading manifest template 'MANIFEST.in'\n",
            "    writing manifest file 'pip-egg-info/torch_spline_conv.egg-info/SOURCES.txt'\n",
            "  Source in /tmp/pip-install-a980v7k7/torch-spline-conv has version 1.1.0, which satisfies requirement torch-spline-conv from https://files.pythonhosted.org/packages/3c/dd/daa9d0b7b2ede913e573876ae286a58ec296678858f2814ff6d6789b234f/torch_spline_conv-1.1.0.tar.gz#sha256=e6029526205d1f7cb535389bebd81decf0649a20ea6a67688c02bd335a7f9339\n",
            "  Removed torch-spline-conv from https://files.pythonhosted.org/packages/3c/dd/daa9d0b7b2ede913e573876ae286a58ec296678858f2814ff6d6789b234f/torch_spline_conv-1.1.0.tar.gz#sha256=e6029526205d1f7cb535389bebd81decf0649a20ea6a67688c02bd335a7f9339 from build tracker '/tmp/pip-req-tracker-8vfpxnhu'\n",
            "Building wheels for collected packages: torch-spline-conv\n",
            "  Created temporary directory: /tmp/pip-wheel-n7ry37hu\n",
            "  Building wheel for torch-spline-conv (setup.py) ... \u001b[?25l  Destination directory: /tmp/pip-wheel-n7ry37hu\n",
            "  Running command /usr/bin/python3 -u -c 'import sys, setuptools, tokenize; sys.argv[0] = '\"'\"'/tmp/pip-install-a980v7k7/torch-spline-conv/setup.py'\"'\"'; __file__='\"'\"'/tmp/pip-install-a980v7k7/torch-spline-conv/setup.py'\"'\"';f=getattr(tokenize, '\"'\"'open'\"'\"', open)(__file__);code=f.read().replace('\"'\"'\\r\\n'\"'\"', '\"'\"'\\n'\"'\"');f.close();exec(compile(code, __file__, '\"'\"'exec'\"'\"'))' bdist_wheel -d /tmp/pip-wheel-n7ry37hu --python-tag cp36\n",
            "  running bdist_wheel\n",
            "  running build\n",
            "  running build_py\n",
            "  creating build\n",
            "  creating build/lib.linux-x86_64-3.6\n",
            "  creating build/lib.linux-x86_64-3.6/test\n",
            "  copying test/utils.py -> build/lib.linux-x86_64-3.6/test\n",
            "  copying test/__init__.py -> build/lib.linux-x86_64-3.6/test\n",
            "  copying test/test_conv.py -> build/lib.linux-x86_64-3.6/test\n",
            "  copying test/test_weighting.py -> build/lib.linux-x86_64-3.6/test\n",
            "  copying test/test_basis.py -> build/lib.linux-x86_64-3.6/test\n",
            "  creating build/lib.linux-x86_64-3.6/torch_spline_conv\n",
            "  copying torch_spline_conv/__init__.py -> build/lib.linux-x86_64-3.6/torch_spline_conv\n",
            "  copying torch_spline_conv/basis.py -> build/lib.linux-x86_64-3.6/torch_spline_conv\n",
            "  copying torch_spline_conv/conv.py -> build/lib.linux-x86_64-3.6/torch_spline_conv\n",
            "  copying torch_spline_conv/weighting.py -> build/lib.linux-x86_64-3.6/torch_spline_conv\n",
            "  creating build/lib.linux-x86_64-3.6/torch_spline_conv/utils\n",
            "  copying torch_spline_conv/utils/__init__.py -> build/lib.linux-x86_64-3.6/torch_spline_conv/utils\n",
            "  copying torch_spline_conv/utils/degree.py -> build/lib.linux-x86_64-3.6/torch_spline_conv/utils\n",
            "  running build_ext\n",
            "  building 'torch_spline_conv.basis_cpu' extension\n",
            "  creating build/temp.linux-x86_64-3.6\n",
            "  creating build/temp.linux-x86_64-3.6/cpu\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/include/python3.6m -c cpu/basis.cpp -o build/temp.linux-x86_64-3.6/cpu/basis.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=basis_cpu -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cpu/basis.o -o build/lib.linux-x86_64-3.6/torch_spline_conv/basis_cpu.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_spline_conv.weighting_cpu' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/include/python3.6m -c cpu/weighting.cpp -o build/temp.linux-x86_64-3.6/cpu/weighting.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=weighting_cpu -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cpu/weighting.o -o build/lib.linux-x86_64-3.6/torch_spline_conv/weighting_cpu.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_spline_conv.basis_cuda' extension\n",
            "  creating build/temp.linux-x86_64-3.6/cuda\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/basis.cpp -o build/temp.linux-x86_64-3.6/cuda/basis.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=basis_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/basis_kernel.cu -o build/temp.linux-x86_64-3.6/cuda/basis_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=basis_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cuda/basis.o build/temp.linux-x86_64-3.6/cuda/basis_kernel.o -L/usr/local/cuda/lib64 -lcudart -o build/lib.linux-x86_64-3.6/torch_spline_conv/basis_cuda.cpython-36m-x86_64-linux-gnu.so\n",
            "  building 'torch_spline_conv.weighting_cuda' extension\n",
            "  x86_64-linux-gnu-gcc -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/weighting.cpp -o build/temp.linux-x86_64-3.6/cuda/weighting.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=weighting_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c cuda/weighting_kernel.cu -o build/temp.linux-x86_64-3.6/cuda/weighting_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=weighting_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11\n",
            "  x86_64-linux-gnu-g++ -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 build/temp.linux-x86_64-3.6/cuda/weighting.o build/temp.linux-x86_64-3.6/cuda/weighting_kernel.o -L/usr/local/cuda/lib64 -lcudart -o build/lib.linux-x86_64-3.6/torch_spline_conv/weighting_cuda.cpython-36m-x86_64-linux-gnu.so\n",
            "  installing to build/bdist.linux-x86_64/wheel\n",
            "  running install\n",
            "  running install_lib\n",
            "  creating build/bdist.linux-x86_64\n",
            "  creating build/bdist.linux-x86_64/wheel\n",
            "  creating build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/utils.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/__init__.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_conv.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_weighting.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  copying build/lib.linux-x86_64-3.6/test/test_basis.py -> build/bdist.linux-x86_64/wheel/test\n",
            "  creating build/bdist.linux-x86_64/wheel/torch_spline_conv\n",
            "  copying build/lib.linux-x86_64-3.6/torch_spline_conv/__init__.py -> build/bdist.linux-x86_64/wheel/torch_spline_conv\n",
            "  copying build/lib.linux-x86_64-3.6/torch_spline_conv/basis.py -> build/bdist.linux-x86_64/wheel/torch_spline_conv\n",
            "  copying build/lib.linux-x86_64-3.6/torch_spline_conv/conv.py -> build/bdist.linux-x86_64/wheel/torch_spline_conv\n",
            "  copying build/lib.linux-x86_64-3.6/torch_spline_conv/basis_cuda.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_spline_conv\n",
            "  creating build/bdist.linux-x86_64/wheel/torch_spline_conv/utils\n",
            "  copying build/lib.linux-x86_64-3.6/torch_spline_conv/utils/__init__.py -> build/bdist.linux-x86_64/wheel/torch_spline_conv/utils\n",
            "  copying build/lib.linux-x86_64-3.6/torch_spline_conv/utils/degree.py -> build/bdist.linux-x86_64/wheel/torch_spline_conv/utils\n",
            "  copying build/lib.linux-x86_64-3.6/torch_spline_conv/weighting.py -> build/bdist.linux-x86_64/wheel/torch_spline_conv\n",
            "  copying build/lib.linux-x86_64-3.6/torch_spline_conv/basis_cpu.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_spline_conv\n",
            "  copying build/lib.linux-x86_64-3.6/torch_spline_conv/weighting_cuda.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_spline_conv\n",
            "  copying build/lib.linux-x86_64-3.6/torch_spline_conv/weighting_cpu.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/torch_spline_conv\n",
            "  running install_egg_info\n",
            "  running egg_info\n",
            "  writing torch_spline_conv.egg-info/PKG-INFO\n",
            "  writing dependency_links to torch_spline_conv.egg-info/dependency_links.txt\n",
            "  writing top-level names to torch_spline_conv.egg-info/top_level.txt\n",
            "  reading manifest file 'torch_spline_conv.egg-info/SOURCES.txt'\n",
            "  reading manifest template 'MANIFEST.in'\n",
            "  writing manifest file 'torch_spline_conv.egg-info/SOURCES.txt'\n",
            "  Copying torch_spline_conv.egg-info to build/bdist.linux-x86_64/wheel/torch_spline_conv-1.1.0-py3.6.egg-info\n",
            "  running install_scripts\n",
            "  adding license file \"LICENSE\" (matched pattern \"LICEN[CS]E*\")\n",
            "  creating build/bdist.linux-x86_64/wheel/torch_spline_conv-1.1.0.dist-info/WHEEL\n",
            "  creating '/tmp/pip-wheel-n7ry37hu/torch_spline_conv-1.1.0-cp36-cp36m-linux_x86_64.whl' and adding 'build/bdist.linux-x86_64/wheel' to it\n",
            "  adding 'test/__init__.py'\n",
            "  adding 'test/test_basis.py'\n",
            "  adding 'test/test_conv.py'\n",
            "  adding 'test/test_weighting.py'\n",
            "  adding 'test/utils.py'\n",
            "  adding 'torch_spline_conv/__init__.py'\n",
            "  adding 'torch_spline_conv/basis.py'\n",
            "  adding 'torch_spline_conv/basis_cpu.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_spline_conv/basis_cuda.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_spline_conv/conv.py'\n",
            "  adding 'torch_spline_conv/weighting.py'\n",
            "  adding 'torch_spline_conv/weighting_cpu.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_spline_conv/weighting_cuda.cpython-36m-x86_64-linux-gnu.so'\n",
            "  adding 'torch_spline_conv/utils/__init__.py'\n",
            "  adding 'torch_spline_conv/utils/degree.py'\n",
            "  adding 'torch_spline_conv-1.1.0.dist-info/LICENSE'\n",
            "  adding 'torch_spline_conv-1.1.0.dist-info/METADATA'\n",
            "  adding 'torch_spline_conv-1.1.0.dist-info/WHEEL'\n",
            "  adding 'torch_spline_conv-1.1.0.dist-info/top_level.txt'\n",
            "  adding 'torch_spline_conv-1.1.0.dist-info/RECORD'\n",
            "  removing build/bdist.linux-x86_64/wheel\n",
            "\u001b[?25hdone\n",
            "  Created wheel for torch-spline-conv: filename=torch_spline_conv-1.1.0-cp36-cp36m-linux_x86_64.whl size=5263533 sha256=57fdc57973a7abf0b677306c5c9f84ccfe1002072e97bd8ea9c1bb5c0e154f57\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-mzwr_9wr/wheels/c8/00/1f/9c414a9a5f340dd8d2e5e362d1bb5cad91fc7aec78c935fd66\n",
            "  Removing source in /tmp/pip-install-a980v7k7/torch-spline-conv\n",
            "Successfully built torch-spline-conv\n",
            "Installing collected packages: torch-spline-conv\n",
            "\n",
            "Successfully installed torch-spline-conv-1.1.0\n",
            "Cleaning up...\n",
            "Removed build tracker '/tmp/pip-req-tracker-8vfpxnhu'\n",
            "Requirement already satisfied: torch-geometric in /usr/local/lib/python3.6/dist-packages (1.3.1)\n",
            "Requirement already satisfied: rdflib in /usr/local/lib/python3.6/dist-packages (from torch-geometric) (4.2.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch-geometric) (1.16.5)\n",
            "Requirement already satisfied: plyfile in /usr/local/lib/python3.6/dist-packages (from torch-geometric) (0.7)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.6/dist-packages (from torch-geometric) (2.3)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.6/dist-packages (from torch-geometric) (0.24.2)\n",
            "Requirement already satisfied: h5py in /usr/local/lib/python3.6/dist-packages (from torch-geometric) (2.8.0)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from torch-geometric) (1.3.1)\n",
            "Requirement already satisfied: googledrivedownloader in /usr/local/lib/python3.6/dist-packages (from torch-geometric) (0.4)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.6/dist-packages (from torch-geometric) (0.21.3)\n",
            "Requirement already satisfied: isodate in /usr/local/lib/python3.6/dist-packages (from rdflib->torch-geometric) (0.6.0)\n",
            "Requirement already satisfied: pyparsing in /usr/local/lib/python3.6/dist-packages (from rdflib->torch-geometric) (2.4.2)\n",
            "Requirement already satisfied: decorator>=4.3.0 in /usr/local/lib/python3.6/dist-packages (from networkx->torch-geometric) (4.4.0)\n",
            "Requirement already satisfied: python-dateutil>=2.5.0 in /usr/local/lib/python3.6/dist-packages (from pandas->torch-geometric) (2.5.3)\n",
            "Requirement already satisfied: pytz>=2011k in /usr/local/lib/python3.6/dist-packages (from pandas->torch-geometric) (2018.9)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from h5py->torch-geometric) (1.12.0)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.6/dist-packages (from scikit-learn->torch-geometric) (0.13.2)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FSgIGVpGnAmy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install syft"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FL6JZ1KQmtWE",
        "colab_type": "text"
      },
      "source": [
        "This class describes all the hyper-parameters for the training. Note that they are all public here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KGFEvtzZmtWG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Arguments():\n",
        "    def __init__(self):\n",
        "        self.batch_size = 64\n",
        "        self.test_batch_size = 64\n",
        "        self.epochs = 10\n",
        "        self.lr = 0.02\n",
        "        self.seed = 1\n",
        "        self.log_interval = 1 # Log info at each batch\n",
        "        self.precision_fractional = 3\n",
        "\n",
        "args = Arguments()\n",
        "\n",
        "_ = torch.manual_seed(args.seed)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZSvuHw4LmtWN",
        "colab_type": "text"
      },
      "source": [
        "Here are PySyft imports. We connect to two remote workers that be call `alice` and `bob` and request another worker called the `crypto_provider` who gives all the crypto primitives we may need."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s98w7mYomtWQ",
        "colab_type": "code",
        "outputId": "4544ec2e-5035-4e36-d673-7625ed1c77bf",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "import syft as sy  # import the Pysyft library\n",
        "hook = sy.TorchHook(torch)  # hook PyTorch to add extra functionalities like Federated and Encrypted Learning\n",
        "\n",
        "# simulation functions\n",
        "def connect_to_workers(n_workers):\n",
        "    return [\n",
        "        sy.VirtualWorker(hook, id=f\"worker{i+1}\")\n",
        "        for i in range(n_workers)\n",
        "    ]\n",
        "def connect_to_crypto_provider():\n",
        "    return sy.VirtualWorker(hook, id=\"crypto_provider\")\n",
        "\n",
        "workers = connect_to_workers(n_workers=2)\n",
        "crypto_provider = connect_to_crypto_provider()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING:root:Torch was already hooked... skipping hooking process\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d_71YeEKmtWb",
        "colab_type": "text"
      },
      "source": [
        "## Getting access and secret share data\n",
        "\n",
        "Here we're using a utility function which simulates the following behaviour: we assume the dataset is distributed in parts each of which is held by one of our workers. The workers then split their data in batches and secret share their data between each others. The final object returned is an iterable on these secret shared batches, that we call the **private data loader**. Note that during the process the local worker (so us) never had access to the data.\n",
        "\n",
        "We obtain as usual a training and testing private dataset, and both the inputs and labels are secret shared."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VFD3gaWemtWe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from torch_geometric.datasets import PPI\n",
        "from torch_geometric.data import DataLoader\n",
        "from torch_geometric.nn import GATConv\n",
        "\n",
        "# We don't use the whole dataset for efficiency purpose, but feel free to increase these numbers\n",
        "n_train_items = 640\n",
        "n_test_items = 640\n",
        "\n",
        "path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'PPI')\n",
        "train_dataset = PPI(path, split='train')\n",
        "val_dataset = PPI(path, split='val')\n",
        "test_dataset = PPI(path, split='test')\n",
        "# val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)\n",
        "\n",
        "\n",
        "def get_private_data_loaders(precision_fractional, workers, crypto_provider):\n",
        "    \n",
        "    def one_hot_of(index_tensor):\n",
        "        \"\"\"\n",
        "        Transform to one hot tensor\n",
        "        \n",
        "        Example:\n",
        "            [0, 3, 9]\n",
        "            =>\n",
        "            [[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
        "             [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],\n",
        "             [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]]\n",
        "            \n",
        "        \"\"\"\n",
        "        onehot_tensor = torch.zeros(*index_tensor.shape, 10) # 10 classes for MNIST\n",
        "        onehot_tensor = onehot_tensor.scatter(1, index_tensor.view(-1, 1), 1)\n",
        "        return onehot_tensor\n",
        "        \n",
        "    def secret_share(tensor):\n",
        "        \"\"\"\n",
        "        Transform to fixed precision and secret share a tensor\n",
        "        \"\"\"\n",
        "        return (\n",
        "            tensor\n",
        "            .fix_precision(precision_fractional=precision_fractional)\n",
        "            .share(*workers, crypto_provider=crypto_provider, requires_grad=True)\n",
        "        )\n",
        "    \n",
        "    transformation = transforms.Compose([\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize((0.1307,), (0.3081,))\n",
        "    ])\n",
        "    \n",
        "    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)\n",
        "    \n",
        "    private_train_loader = [\n",
        "        (secret_share(data), secret_share(one_hot_of(target)))\n",
        "        for i, (data, target) in enumerate(train_loader)\n",
        "        if i < n_train_items / args.batch_size\n",
        "    ]\n",
        "    \n",
        "    test_loader = DataLoader(test_dataset, batch_size=args.test_batch_size, shuffle=False)\n",
        "    \n",
        "    private_test_loader = [\n",
        "        (secret_share(data), secret_share(target.float()))\n",
        "        for i, (data, target) in enumerate(test_loader)\n",
        "        if i < n_test_items / args.test_batch_size\n",
        "    ]\n",
        "    \n",
        "    return private_train_loader, private_test_loader\n",
        "    \n",
        "    \n",
        "private_train_loader, private_test_loader = get_private_data_loaders(\n",
        "    precision_fractional=args.precision_fractional,\n",
        "    workers=workers,\n",
        "    crypto_provider=crypto_provider\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z6V2uEeXmtWn",
        "colab_type": "text"
      },
      "source": [
        "## Model specification\n",
        "\n",
        "Here is the model that we will use, it's a rather simple one but [it has proved to perform reasonably well on MNIST](https://towardsdatascience.com/handwritten-digit-mnist-pytorch-977b5338e627)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JIA9Ix7imtWq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Net(torch.nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.conv1 = GATConv(train_dataset.num_features, 256, heads=4)\n",
        "        self.lin1 = nn.Linear(train_dataset.num_features, 4 * 256)\n",
        "        self.conv2 = GATConv(4 * 256, 256, heads=4)\n",
        "        self.lin2 = nn.Linear(4 * 256, 4 * 256)\n",
        "        self.conv3 = GATConv(4 * 256, train_dataset.num_classes, heads=6, concat=False)\n",
        "        self.lin3 = nn.Linear(4 * 256, train_dataset.num_classes)\n",
        "\n",
        "    def forward(self, x, edge_index):\n",
        "        x = F.elu(self.conv1(x, edge_index) + self.lin1(x))\n",
        "        x = F.elu(self.conv2(x, edge_index) + self.lin2(x))\n",
        "        x = self.conv3(x, edge_index) + self.lin3(x)\n",
        "        return x\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l-n6W8ANmtW0",
        "colab_type": "text"
      },
      "source": [
        "## Training and testing functions\n",
        "\n",
        "The training is done almost as usual, the real difference is that we can't use losses like negative log-likelihood (`F.nll_loss` in PyTorch) because it's quite complicated to reproduce these functions with SMPC. Instead, we use a simpler Mean Square Error loss."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P66OH8d7mtW3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(args, model, private_train_loader, loss_op, optimizer, epoch):\n",
        "    model.train()\n",
        "    for batch_idx, (data, target) in enumerate(private_train_loader): # <-- now it is a private dataset\n",
        "        start_time = time.time()\n",
        "        num_graphs = data.num_graphs\n",
        "        data.batch = None        \n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        output = model(data)\n",
        "        \n",
        "        # loss = F.nll_loss(output, target)  <-- not possible here\n",
        "        batch_size = output.shape[0]\n",
        "        loss = ((output - target)**2).sum().refresh()/batch_size\n",
        "        \n",
        "        loss = loss_op(model(data.x, data.edge_index), data.y)\n",
        "        total_loss += loss.item() * num_graphs\n",
        "\n",
        "        loss.backward()\n",
        "        \n",
        "        optimizer.step()\n",
        "\n",
        "        if batch_idx % args.log_interval == 0:\n",
        "            loss = loss.get().float_precision()\n",
        "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}\\tTime: {:.3f}s'.format(\n",
        "                epoch, batch_idx * args.batch_size, len(private_train_loader) * args.batch_size,\n",
        "                100. * batch_idx / len(private_train_loader), total_loss / len(train_loader.dataset), time.time() - start_time))\n",
        "            "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wu11u-ICmtXE",
        "colab_type": "text"
      },
      "source": [
        "The test function "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JITe2x1lmtXI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def test(args, model, private_test_loader):\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    with torch.no_grad():\n",
        "        for data, target in private_test_loader:\n",
        "            start_time = time.time()\n",
        "            \n",
        "            output = model(data)\n",
        "            pred = output.argmax(dim=1)\n",
        "            correct += pred.eq(target.view_as(pred)).sum()\n",
        "\n",
        "    correct = correct.get().float_precision()\n",
        "    print('\\nTest set: Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
        "        correct.item(), len(private_test_loader)* args.test_batch_size,\n",
        "        100. * correct.item() / (len(private_test_loader) * args.test_batch_size)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZI2ZDy9OmtXt",
        "colab_type": "text"
      },
      "source": [
        "### Let's launch the training !\n",
        "\n",
        "A few notes about what's happening here. First, we secret share all the model parameters across our workers. Second, we convert optimizer's hyperparameters to fixed precision. Note that we don't need to secret share them because they are public in our context, but as secret shared values live in finite fields we still need to move them in finite fields using using `.fix_precision`, in order to perform consistently operations like the weight update $W \\leftarrow W - \\alpha * \\Delta W$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": true,
        "id": "HYRj-kqFmtXy",
        "colab_type": "code",
        "outputId": "1b6692d7-db19-40d3-e046-69e10ea819e5",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "model = Net()\n",
        "model = model.fix_precision().share(*workers, crypto_provider=crypto_provider, requires_grad=True)\n",
        "\n",
        "loss_op = nn.BCEWithLogitsLoss()\n",
        "loss_op = loss_op.fix_precision()\n",
        "\n",
        "optimizer = optim.Adam(model.parameters(), lr=0.005)\n",
        "optimizer = optimizer.fix_precision() \n",
        "\n",
        "for epoch in range(1, args.epochs + 1):\n",
        "    train(args, model, private_train_loader, optimizer, epoch)\n",
        "    test(args, model, private_test_loader)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Train Epoch: 1 [0/640 (0%)]\tLoss: 1.128000\tTime: 4.995s\n",
            "Train Epoch: 1 [64/640 (10%)]\tLoss: 1.011000\tTime: 5.022s\n",
            "Train Epoch: 1 [128/640 (20%)]\tLoss: 0.990000\tTime: 5.024s\n",
            "Train Epoch: 1 [192/640 (30%)]\tLoss: 0.902000\tTime: 4.977s\n",
            "Train Epoch: 1 [256/640 (40%)]\tLoss: 0.887000\tTime: 4.926s\n",
            "Train Epoch: 1 [320/640 (50%)]\tLoss: 0.875000\tTime: 4.924s\n",
            "Train Epoch: 1 [384/640 (60%)]\tLoss: 0.853000\tTime: 4.878s\n",
            "Train Epoch: 1 [448/640 (70%)]\tLoss: 0.849000\tTime: 4.870s\n",
            "Train Epoch: 1 [512/640 (80%)]\tLoss: 0.830000\tTime: 4.851s\n",
            "Train Epoch: 1 [576/640 (90%)]\tLoss: 0.839000\tTime: 4.856s\n",
            "\n",
            "Test set: Accuracy: 300.0/640 (47%)\n",
            "\n",
            "Train Epoch: 2 [0/640 (0%)]\tLoss: 0.782000\tTime: 4.846s\n",
            "Train Epoch: 2 [64/640 (10%)]\tLoss: 0.732000\tTime: 4.825s\n",
            "Train Epoch: 2 [128/640 (20%)]\tLoss: 0.794000\tTime: 4.798s\n",
            "Train Epoch: 2 [192/640 (30%)]\tLoss: 0.717000\tTime: 4.810s\n",
            "Train Epoch: 2 [256/640 (40%)]\tLoss: 0.705000\tTime: 4.862s\n",
            "Train Epoch: 2 [320/640 (50%)]\tLoss: 0.707000\tTime: 4.819s\n",
            "Train Epoch: 2 [384/640 (60%)]\tLoss: 0.703000\tTime: 4.823s\n",
            "Train Epoch: 2 [448/640 (70%)]\tLoss: 0.720000\tTime: 4.809s\n",
            "Train Epoch: 2 [512/640 (80%)]\tLoss: 0.711000\tTime: 4.852s\n",
            "Train Epoch: 2 [576/640 (90%)]\tLoss: 0.745000\tTime: 4.943s\n",
            "\n",
            "Test set: Accuracy: 462.0/640 (72%)\n",
            "\n",
            "Train Epoch: 3 [0/640 (0%)]\tLoss: 0.668000\tTime: 4.801s\n",
            "Train Epoch: 3 [64/640 (10%)]\tLoss: 0.599000\tTime: 4.811s\n",
            "Train Epoch: 3 [128/640 (20%)]\tLoss: 0.699000\tTime: 4.810s\n",
            "Train Epoch: 3 [192/640 (30%)]\tLoss: 0.601000\tTime: 4.799s\n",
            "Train Epoch: 3 [256/640 (40%)]\tLoss: 0.591000\tTime: 5.003s\n",
            "Train Epoch: 3 [320/640 (50%)]\tLoss: 0.592000\tTime: 4.806s\n",
            "Train Epoch: 3 [384/640 (60%)]\tLoss: 0.603000\tTime: 4.800s\n",
            "Train Epoch: 3 [448/640 (70%)]\tLoss: 0.629000\tTime: 4.838s\n",
            "Train Epoch: 3 [512/640 (80%)]\tLoss: 0.625000\tTime: 4.864s\n",
            "Train Epoch: 3 [576/640 (90%)]\tLoss: 0.669000\tTime: 4.864s\n",
            "\n",
            "Test set: Accuracy: 497.0/640 (78%)\n",
            "\n",
            "Train Epoch: 4 [0/640 (0%)]\tLoss: 0.582000\tTime: 4.806s\n",
            "Train Epoch: 4 [64/640 (10%)]\tLoss: 0.503000\tTime: 4.845s\n",
            "Train Epoch: 4 [128/640 (20%)]\tLoss: 0.624000\tTime: 4.817s\n",
            "Train Epoch: 4 [192/640 (30%)]\tLoss: 0.516000\tTime: 4.804s\n",
            "Train Epoch: 4 [256/640 (40%)]\tLoss: 0.514000\tTime: 4.824s\n",
            "Train Epoch: 4 [320/640 (50%)]\tLoss: 0.511000\tTime: 4.804s\n",
            "Train Epoch: 4 [384/640 (60%)]\tLoss: 0.532000\tTime: 4.876s\n",
            "Train Epoch: 4 [448/640 (70%)]\tLoss: 0.564000\tTime: 4.874s\n",
            "Train Epoch: 4 [512/640 (80%)]\tLoss: 0.559000\tTime: 4.861s\n",
            "Train Epoch: 4 [576/640 (90%)]\tLoss: 0.618000\tTime: 4.813s\n",
            "\n",
            "Test set: Accuracy: 517.0/640 (81%)\n",
            "\n",
            "Train Epoch: 5 [0/640 (0%)]\tLoss: 0.525000\tTime: 4.823s\n",
            "Train Epoch: 5 [64/640 (10%)]\tLoss: 0.441000\tTime: 4.831s\n",
            "Train Epoch: 5 [128/640 (20%)]\tLoss: 0.571000\tTime: 4.857s\n",
            "Train Epoch: 5 [192/640 (30%)]\tLoss: 0.458000\tTime: 4.822s\n",
            "Train Epoch: 5 [256/640 (40%)]\tLoss: 0.460000\tTime: 4.806s\n",
            "Train Epoch: 5 [320/640 (50%)]\tLoss: 0.453000\tTime: 4.834s\n",
            "Train Epoch: 5 [384/640 (60%)]\tLoss: 0.480000\tTime: 4.913s\n",
            "Train Epoch: 5 [448/640 (70%)]\tLoss: 0.515000\tTime: 4.846s\n",
            "Train Epoch: 5 [512/640 (80%)]\tLoss: 0.507000\tTime: 4.839s\n",
            "Train Epoch: 5 [576/640 (90%)]\tLoss: 0.576000\tTime: 4.815s\n",
            "\n",
            "Test set: Accuracy: 529.0/640 (83%)\n",
            "\n",
            "Train Epoch: 6 [0/640 (0%)]\tLoss: 0.475000\tTime: 4.808s\n",
            "Train Epoch: 6 [64/640 (10%)]\tLoss: 0.392000\tTime: 4.811s\n",
            "Train Epoch: 6 [128/640 (20%)]\tLoss: 0.528000\tTime: 4.863s\n",
            "Train Epoch: 6 [192/640 (30%)]\tLoss: 0.410000\tTime: 4.805s\n",
            "Train Epoch: 6 [256/640 (40%)]\tLoss: 0.414000\tTime: 4.860s\n",
            "Train Epoch: 6 [320/640 (50%)]\tLoss: 0.408000\tTime: 4.853s\n",
            "Train Epoch: 6 [384/640 (60%)]\tLoss: 0.438000\tTime: 4.911s\n",
            "Train Epoch: 6 [448/640 (70%)]\tLoss: 0.474000\tTime: 4.879s\n",
            "Train Epoch: 6 [512/640 (80%)]\tLoss: 0.466000\tTime: 4.837s\n",
            "Train Epoch: 6 [576/640 (90%)]\tLoss: 0.543000\tTime: 4.805s\n",
            "\n",
            "Test set: Accuracy: 539.0/640 (84%)\n",
            "\n",
            "Train Epoch: 7 [0/640 (0%)]\tLoss: 0.437000\tTime: 4.789s\n",
            "Train Epoch: 7 [64/640 (10%)]\tLoss: 0.354000\tTime: 4.804s\n",
            "Train Epoch: 7 [128/640 (20%)]\tLoss: 0.489000\tTime: 4.830s\n",
            "Train Epoch: 7 [192/640 (30%)]\tLoss: 0.374000\tTime: 4.870s\n",
            "Train Epoch: 7 [256/640 (40%)]\tLoss: 0.379000\tTime: 4.857s\n",
            "Train Epoch: 7 [320/640 (50%)]\tLoss: 0.368000\tTime: 4.831s\n",
            "Train Epoch: 7 [384/640 (60%)]\tLoss: 0.403000\tTime: 4.852s\n",
            "Train Epoch: 7 [448/640 (70%)]\tLoss: 0.442000\tTime: 4.815s\n",
            "Train Epoch: 7 [512/640 (80%)]\tLoss: 0.432000\tTime: 4.814s\n",
            "Train Epoch: 7 [576/640 (90%)]\tLoss: 0.508000\tTime: 4.852s\n",
            "\n",
            "Test set: Accuracy: 544.0/640 (85%)\n",
            "\n",
            "Train Epoch: 8 [0/640 (0%)]\tLoss: 0.405000\tTime: 4.804s\n",
            "Train Epoch: 8 [64/640 (10%)]\tLoss: 0.324000\tTime: 4.821s\n",
            "Train Epoch: 8 [128/640 (20%)]\tLoss: 0.458000\tTime: 4.830s\n",
            "Train Epoch: 8 [192/640 (30%)]\tLoss: 0.347000\tTime: 4.880s\n",
            "Train Epoch: 8 [256/640 (40%)]\tLoss: 0.352000\tTime: 4.858s\n",
            "Train Epoch: 8 [320/640 (50%)]\tLoss: 0.341000\tTime: 4.901s\n",
            "Train Epoch: 8 [384/640 (60%)]\tLoss: 0.374000\tTime: 4.923s\n",
            "Train Epoch: 8 [448/640 (70%)]\tLoss: 0.416000\tTime: 4.873s\n",
            "Train Epoch: 8 [512/640 (80%)]\tLoss: 0.402000\tTime: 4.818s\n",
            "Train Epoch: 8 [576/640 (90%)]\tLoss: 0.483000\tTime: 4.816s\n",
            "\n",
            "Test set: Accuracy: 553.0/640 (86%)\n",
            "\n",
            "Train Epoch: 9 [0/640 (0%)]\tLoss: 0.382000\tTime: 4.831s\n",
            "Train Epoch: 9 [64/640 (10%)]\tLoss: 0.303000\tTime: 4.864s\n",
            "Train Epoch: 9 [128/640 (20%)]\tLoss: 0.433000\tTime: 4.857s\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}