does-chess-board-have-boundary-tilting.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## This ipynb checkpoint contains the experiment in Appendix C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader , TensorDataset\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.cluster import SpectralClustering\n",
    "import sklearn.datasets\n",
    "from tqdm import tqdm, trange\n",
    "\n",
    "from sklearn.metrics import accuracy_score\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "seed = 1\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "block = 9\n",
    "num_points_each_tile = 100\n",
    "dim = 100\n",
    "num_non_robust_features = 1\n",
    "non_robust_feature_width = 0.05\n",
    "#non_robust_feature_width = 0\n",
    "noise_sigma = 0.01\n",
    "#noise_sigma = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_original = np.zeros((block**2,2))\n",
    "y_original = np.zeros(block**2)\n",
    "\n",
    "tmp = np.arange(block)\n",
    "x_original[:,0] = np.tile(tmp,block)\n",
    "x_original[:,1] = np.repeat(tmp,block)\n",
    "y_original[:] = np.mod(x_original[:,0] + x_original[:,1],2)\n",
    "\n",
    "x = np.zeros((block**2*num_points_each_tile, 2))\n",
    "x[:,0] = np.repeat(x_original[:,0],num_points_each_tile)\n",
    "x[:,1] = np.repeat(x_original[:,1],num_points_each_tile)\n",
    "y = np.repeat(y_original,num_points_each_tile)\n",
    "y = y.astype(int)\n",
    "\n",
    "num_points_each_edge = num_points_each_tile//4\n",
    "\n",
    "noise = np.zeros((0,2))\n",
    "\n",
    "margin = 0.3\n",
    "width = 0.2\n",
    "\n",
    "for _ in range(block**2):\n",
    "\n",
    "    edge1 = np.zeros((num_points_each_edge,2))\n",
    "    edge1[:,0] = (np.random.rand(num_points_each_edge) - 0.5)* (1 - 2*margin)\n",
    "    edge1[:,1] = np.random.rand(num_points_each_edge)*width + (0.5-width-margin)\n",
    "\n",
    "    edge2 = np.zeros((num_points_each_edge,2))\n",
    "    edge2[:,0] = (np.random.rand(num_points_each_edge) - 0.5)* (1 - 2*margin)\n",
    "    edge2[:,1] = -np.random.rand(num_points_each_edge)*width - (0.5-width-margin)\n",
    "\n",
    "    edge3 = np.zeros((num_points_each_edge,2))\n",
    "    edge3[:,1] = (np.random.rand(num_points_each_edge) - 0.5)* (1 - 2*margin)\n",
    "    edge3[:,0] = np.random.rand(num_points_each_edge)*width + (0.5-width-margin)\n",
    "\n",
    "    edge4 = np.zeros((num_points_each_edge,2))\n",
    "    edge4[:,1] = (np.random.rand(num_points_each_edge) - 0.5)* (1 - 2*margin)\n",
    "    edge4[:,0] = -np.random.rand(num_points_each_edge)*width - (0.5-width-margin)\n",
    "\n",
    "    edges = np.vstack([edge1, edge2, edge3, edge4])\n",
    "\n",
    "    noise = np.vstack([noise,edges])\n",
    "\n",
    "x = x + noise\n",
    "x -= (block-1)/2\n",
    "\n",
    "xx = np.zeros((block**2*num_points_each_tile, dim))\n",
    "xx[:,:2] = x\n",
    "\n",
    "# Add non-robust feature on the third dimension\n",
    "\n",
    "for j in range(2, 2 + num_non_robust_features):\n",
    "\n",
    "    non_robust_blocks = np.random.choice([0, 1], size=(block**2), p=[1./2, 1./2])\n",
    "    xx[:,j] = non_robust_feature_width * np.repeat(non_robust_blocks, num_points_each_tile)\n",
    "\n",
    "# Add noise on the third and higher dimension\n",
    "\n",
    "xx[:,2:] += noise_sigma * np.random.randn(block**2*num_points_each_tile, dim-2)\n",
    "\n",
    "np.save('../data/x.npy',x) \n",
    "np.save('../data/xx.npy',xx)\n",
    "np.save('../data/y.npy',y)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize the 2D-data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8,8))\n",
    "plt.scatter(x[:,0],x[:,1],s=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.load('../data/x.npy').astype(np.float64) \n",
    "xx = np.load('../data/xx.npy').astype(np.float64) \n",
    "y = np.load('../data/y.npy')\n",
    "\n",
    "print( np.min(xx[:,:1]) , np.max(xx[:,:1]) , np.mean(xx[:,:1]),np.std(xx[:,:1]))\n",
    "print( np.min(xx[:,2:]) , np.max(xx[:,2:]) , np.mean(xx[:,2:]),np.std(xx[:,2:]) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "x_train,x_test,y_train,y_test = train_test_split(xx,y,test_size=0.4,random_state=0)\n",
    "\n",
    "print(x_train.shape , y_train.shape)\n",
    "\n",
    "plt.figure(figsize=(15,7))\n",
    "plt.subplot(121)\n",
    "plt.scatter(x_train[:,0],x_train[:,1],s=40,c=y_train,cmap=plt.cm.Spectral)\n",
    "plt.title(\"Training data\")\n",
    "plt.subplot(122)\n",
    "plt.scatter(x_test[:,0],x_test[:,1],s=40,c=y_test,cmap=plt.cm.Spectral)\n",
    "plt.title(\"Test data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_width = 256\n",
    "net_width_fc_layer = 64\n",
    "\n",
    "#net_width = 64\n",
    "#net_width_fc_layer = 16\n",
    "\n",
    "class ResNet(nn.Module):\n",
    "    def __init__(self, input_dim = 100 , num_classes = 2):\n",
    "        super(ResNet, self).__init__()\n",
    "        \n",
    "        self.layer0 = nn.Sequential(\n",
    "            nn.Linear(input_dim, net_width),\n",
    "            nn.BatchNorm1d(net_width),\n",
    "            nn.ReLU(),\n",
    "            #nn.Dropout(0.2)\n",
    "        )\n",
    "        #nn.init.xavier_uniform(self.layer0[0].weight, 1)\n",
    "        \n",
    "        self.layer1 = nn.Sequential(\n",
    "            nn.Linear(net_width, net_width),\n",
    "            nn.BatchNorm1d(net_width),\n",
    "            nn.ReLU(),\n",
    "            #nn.Dropout(0.2)\n",
    "        )\n",
    "        #nn.init.xavier_uniform(self.layer1[0].weight, 1)\n",
    "        \n",
    "        self.layer2 = nn.Sequential(\n",
    "            nn.Linear(net_width, net_width),\n",
    "            nn.BatchNorm1d(net_width),\n",
    "            nn.ReLU(),\n",
    "            #nn.Dropout(0.2)\n",
    "        )\n",
    "        #nn.init.xavier_uniform(self.layer2[0].weight, 1)\n",
    "        \n",
    "        self.layer3 = nn.Sequential(\n",
    "            nn.Linear(net_width, net_width),\n",
    "            nn.BatchNorm1d(net_width),\n",
    "            nn.ReLU(),\n",
    "            #nn.Dropout(0.2)\n",
    "        )\n",
    "        #nn.init.xavier_uniform(self.layer3[0].weight, 1)\n",
    "        \n",
    "        self.layer4 = nn.Sequential(\n",
    "            nn.Linear(net_width, net_width),\n",
    "            nn.BatchNorm1d(net_width),\n",
    "            nn.ReLU(),\n",
    "            #nn.Dropout(0.2)\n",
    "        )\n",
    "        #nn.init.xavier_uniform(self.layer4[0].weight, 1)\n",
    "        \n",
    "        self.layer5 = nn.Sequential(\n",
    "            nn.Linear(net_width, net_width),\n",
    "            nn.BatchNorm1d(net_width),\n",
    "            nn.ReLU(),\n",
    "            #nn.Dropout(0.2)\n",
    "        )\n",
    "        #nn.init.xavier_uniform(self.layer5[0].weight, 1)\n",
    "        \n",
    "        self.layer6 = nn.Sequential(\n",
    "            nn.Linear(net_width, net_width),\n",
    "            nn.BatchNorm1d(net_width),\n",
    "            nn.ReLU(),\n",
    "            #nn.Dropout(0.2)\n",
    "        )\n",
    "            \n",
    "        self.layer7 = nn.Sequential(\n",
    "            nn.Linear(net_width, net_width_fc_layer),\n",
    "            nn.BatchNorm1d(net_width_fc_layer)\n",
    "        )\n",
    "        #nn.init.xavier_uniform(self.layer6[0].weight, 1)\n",
    "        \n",
    "        self.layer8 = nn.Sequential(\n",
    "            nn.Linear(net_width_fc_layer, num_classes)\n",
    "        )\n",
    "        #nn.init.xavier_uniform(self.layer7[0].weight, 1)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        outputs = self.layer0(x)\n",
    "        \n",
    "        outputs1 = self.layer1(outputs)\n",
    "        outputs2 = self.layer2(outputs1 + outputs)\n",
    "        \n",
    "        outputs3 = self.layer3(outputs2)\n",
    "        outputs4 = self.layer4(outputs3 + outputs2 )\n",
    "        \n",
    "        outputs5 = self.layer5(outputs4)\n",
    "        outputs6 = self.layer6(outputs5 + outputs4 )\n",
    "        \n",
    "        outputs7 = self.layer7(outputs6)\n",
    "        outputs8 = self.layer8(outputs7)\n",
    "\n",
    "        return outputs8\n",
    "    \n",
    "    def predict(self,x):\n",
    "        srcdata = torch.from_numpy(x).type(torch.FloatTensor).to(device)\n",
    "        outputs = self.forward(srcdata)\n",
    "        _, y_pred = torch.max(outputs.data, 1)\n",
    "        return y_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "X_train = torch.from_numpy(x_train).type(torch.FloatTensor).to(device)\n",
    "Y_train = torch.from_numpy(y_train).to(device)\n",
    "X_test = torch.from_numpy(x_test).type(torch.FloatTensor).to(device)\n",
    "Y_test = torch.from_numpy(y_test).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Generate_data(x_data, y_data ,batch_size = 128) :\n",
    "    \n",
    "    train_data = TensorDataset(x_data, y_data )\n",
    "\n",
    "    trainloader = torch.utils.data.DataLoader(train_data, \n",
    "                                              batch_size=batch_size\n",
    "                                             )\n",
    "    return trainloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 128\n",
    "train_loader = Generate_data( X_train, Y_train, batch_size = batch_size )\n",
    "test_loader  = Generate_data( X_test,  Y_test,  batch_size = batch_size )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(X_train.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_model(model_tst, loader_tst, dataset_name):\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    model_tst.eval()\n",
    "    with torch.no_grad():\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        for images, labels in loader_tst:\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = model_tst(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "            #print(total)\n",
    "\n",
    "        print('Test Accuracy of the model on the ' + dataset_name + ' : {} %'.format(100 * correct / total))\n",
    "        \n",
    "    return 100 * correct / total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def edge(models , X , y , dim =100  ): \n",
    "    models.eval()\n",
    "    # Set min and max values and give it some padding\n",
    "    x_min, x_max = X[:, 0].min() - 0.5, X[:, 0].max() + 0.5\n",
    "    y_min, y_max = X[:, 1].min() - 0.5, X[:, 1].max() + 0.5\n",
    "    h = 0.04\n",
    "    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 201), np.linspace(y_min, y_max, 201))\n",
    "    # Predict the function value for the whole gid\n",
    "    t = np.c_[xx.ravel(), yy.ravel()]\n",
    "    \n",
    "    tt = np.zeros((t.shape[0],dim))\n",
    "    tt[:,:2] = t\n",
    "    \n",
    "    X_test = torch.from_numpy(tt).float().to(device)\n",
    "    outputs = models(X_test)\n",
    "    _, y_pred = torch.max(outputs.data, 1)\n",
    "    y_pred = y_pred.cpu().reshape(xx.shape)\n",
    "    \n",
    "    # Plot the contour and training examples\n",
    "    \n",
    "    plt.contourf(xx, yy, y_pred , cmap=plt.cm.Spectral)\n",
    "    plt.scatter(X[:, 0], X[:, 1], c=y, cmap=plt.cm.binary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize3D(vis_net, x, y, dir1, dir2, dir3, xlim = [-1,1], ylim = [-1,1], zlim = [-1,1]):\n",
    "    \n",
    "    # Normalize the two directions\n",
    "    print('Take three orthogonal directions')\n",
    "    dir1 = dir1/torch.norm(dir1, p = float('2'))\n",
    "    dir2 = dir2/torch.norm(dir2, p = float('2'))\n",
    "    dir3 = dir3/torch.norm(dir3, p = float('2'))\n",
    "    \n",
    "    # Check if the three directions are orthogonal to each other\n",
    "    inner_product1 = torch.abs(torch.dot(dir1.view(-1), dir2.view(-1)))\n",
    "    inner_product2 = torch.abs(torch.dot(dir1.view(-1), dir3.view(-1)))\n",
    "    inner_product3 = torch.abs(torch.dot(dir2.view(-1), dir3.view(-1)))\n",
    "    \n",
    "    check_inner_product1 = (inner_product1<0.01).item()\n",
    "    check_inner_product2 = (inner_product2<0.01).item()\n",
    "    check_inner_product3 = (inner_product3<0.01).item()\n",
    "\n",
    "    assert check_inner_product1, \"The three directions are not orthogonal\"\n",
    "    assert check_inner_product2, \"The three directions are not orthogonal\"\n",
    "    assert check_inner_product3, \"The three directions are not orthogonal\"\n",
    "    \n",
    "    # Generate the visualization and data grid\n",
    "    #lenx, leny, lenz = 51, 51, 51\n",
    "    xx, yy, zz = np.mgrid[xlim[0]:xlim[1]:30j, ylim[0]:ylim[1]:30j, zlim[0]:zlim[1]:30j]\n",
    "\n",
    "    t = np.c_[xx.ravel(), yy.ravel(), zz.ravel()]\n",
    "    vis_grid = torch.from_numpy(t).float().to(device)\n",
    "    dirs_mat = torch.cat([dir1.reshape(1, -1), dir2.reshape(1, -1), dir3.reshape(1, -1)]).to(device)\n",
    "    x_grid = torch.mm(vis_grid, dirs_mat).to('cpu') + x\n",
    "    \n",
    "    print(x_grid[:,:100])\n",
    "        \n",
    "    grid_output = []\n",
    "    grid_loader = torch.utils.data.DataLoader(TensorDataset(x_grid), batch_size=64, shuffle=False, num_workers=2)\n",
    "    \n",
    "    vis_net.eval()\n",
    "    \n",
    "    softmax1 = nn.Softmax()\n",
    "    \n",
    "    for grid_points in tqdm(grid_loader):\n",
    "        \n",
    "        grid_points = grid_points[0].to(device)\n",
    "        grid_ys = vis_net(grid_points)    \n",
    "        grid_ys = softmax1(grid_ys)\n",
    "        grid_ys = grid_ys[:,y].detach().cpu().numpy()\n",
    "        grid_output.append(grid_ys)\n",
    "        \n",
    "        #_, grid_ys = torch.max(grid_ys.data, 1)\n",
    "        #grid_pred = (grid_ys.to('cpu') == y)\n",
    "        #grid_output.append(grid_pred)\n",
    "    \n",
    "    y_pred0 = np.concatenate(grid_output)\n",
    "    print([np.min(y_pred0),np.max(y_pred0)])\n",
    "    #y_pred = y_pred0.reshape(xx.shape)\n",
    "        \n",
    "    # set the colors of each object\n",
    "    #colors = np.empty(y_pred.shape, dtype=object)\n",
    "    #colors[y_pred] = 'red'\n",
    "    \n",
    "    # and plot everything\n",
    "    fig = go.Figure(data=go.Volume(\n",
    "    x=xx.flatten(),\n",
    "    y=yy.flatten(),\n",
    "    z=zz.flatten(),\n",
    "    value=y_pred0.flatten(),\n",
    "    isomin=0,\n",
    "    isomax=1,\n",
    "    opacity=0.1, # needs to be small to see through all surfaces\n",
    "    surface_count=17, # needs to be a large number for good volume rendering\n",
    "    ))\n",
    "    fig.show()\n",
    "    \n",
    "    return fig\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training using stochastic gradient descent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "# model = NeuralNet().to(device)\n",
    "model = ResNet(input_dim = dim).to(device)\n",
    "\n",
    "def update_lr(optimizer, lr):    \n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = lr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_lr = 0.003\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=initial_lr, momentum=0.9, weight_decay = 5e-4)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 100\n",
    "\n",
    "best_acc = 0\n",
    "train = True\n",
    "#train = False\n",
    "\n",
    "\n",
    "if train :\n",
    "    for epoch in range(num_epochs):\n",
    "        \n",
    "        model.train()\n",
    "        \n",
    "        if epoch > 0.75*num_epochs:\n",
    "            lr = initial_lr/100\n",
    "        elif epoch > 0.5*num_epochs:\n",
    "            lr = initial_lr/10\n",
    "        else:\n",
    "            lr = initial_lr\n",
    "\n",
    "        update_lr(optimizer, lr)\n",
    "\n",
    "        for i, (inputs, targets) in enumerate(train_loader):\n",
    "\n",
    "            inputs = inputs.to(device)\n",
    "            targets = targets.to(device)\n",
    "            \n",
    "            outputs = model(inputs)\n",
    "            #print(outputs.type())\n",
    "            loss = criterion(outputs, targets)\n",
    "\n",
    "            # Backward and optimize\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        \n",
    "        if (epoch+1) % 10 == 0:\n",
    "            print ('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))\n",
    "        \n",
    "        if (epoch+1)% 10 == 0:\n",
    "\n",
    "            acc = eval_model(model,test_loader ,'src test data')\n",
    "\n",
    "            if acc > best_acc:\n",
    "                print('Saving model ' + '..')\n",
    "                state = {\n",
    "                    'net': model.state_dict(),\n",
    "                    'acc': acc,\n",
    "                }\n",
    "                if not os.path.isdir('checkpoint'):\n",
    "                    os.mkdir('checkpoint')\n",
    "                torch.save(state, './checkpoint/Chessboard_'+ str(block)+ '_' + str(block)+'.ckpt')\n",
    "                best_acc = acc\n",
    "    print('best acc :' , best_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = torch.load('./checkpoint/Chessboard_'+ str(block)+ '_' + str(block)+'.pth')\n",
    "model.load_state_dict(checkpoint['net'] )\n",
    "print(checkpoint['acc'])\n",
    "\n",
    "'''\n",
    "print(accuracy_score(model.predict(x_train).cpu(),y_train ))\n",
    "print(accuracy_score(model.predict(x_test).cpu(),y_test ))\n",
    "'''\n",
    "\n",
    "eval_model(model,train_loader ,'src train data')\n",
    "eval_model(model,test_loader ,'src test data')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examine the results on the data manifold (in the middle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,16))\n",
    "plt.subplot(211)\n",
    "edge(model , x_train, y_train , dim = dim)\n",
    "plt.subplot(212)\n",
    "edge(model , x_test, y_test , dim = dim)\n",
    "#plt.savefig('edge.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize a 3D decision boundary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xlim = [-5, 5]\n",
    "ylim = [-5, 5]\n",
    "zlim = [-5, 5]\n",
    "\n",
    "dir0 = torch.from_numpy(np.zeros(dim)).float()\n",
    "dir1 = torch.from_numpy(np.zeros(dim)).float()\n",
    "dir2 = torch.from_numpy(np.zeros(dim)).float()\n",
    "dir0[0] = 1\n",
    "dir1[1] = 1\n",
    "dir2[2] = 1\n",
    "\n",
    "fig = visualize3D(model, torch.from_numpy(np.zeros(dim)).float(), 0, dir0, dir1, dir2, xlim, ylim, zlim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_pytorch_p36)",
   "language": "python",
   "name": "conda_pytorch_p36"
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}