unet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "from torch import optim\n",
    "import torchvision\n",
    "from torch.utils.data import DataLoader\n",
    "from skimage import io\n",
    "import numpy as np\n",
    "import gc\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class double_conv_relu(nn.Module):\n",
    "    \n",
    "    def __init__(self, in_channels, out_channels, dropout=False):\n",
    "        super(double_conv_relu, self).__init__()\n",
    "        \n",
    "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)\n",
    "        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)\n",
    "        self.drop = nn.Dropout2d(p=0.2)\n",
    "        self.norm1 = nn.BatchNorm2d(out_channels)\n",
    "        self.norm2 = nn.BatchNorm2d(out_channels)\n",
    "        self.ReLU = nn.ReLU(inplace=True)\n",
    "        self.dropout = dropout\n",
    "    def forward(self, x):\n",
    "        out = self.conv1(x)\n",
    "        out = self.norm1(out)\n",
    "        out = self.ReLU(out)\n",
    "        out = self.conv2(out)\n",
    "        out = self.norm2(out)\n",
    "        out = self.ReLU(out)\n",
    "        if(self.dropout):\n",
    "            out = self.drop(out)\n",
    "        return out\n",
    "    \n",
    "\n",
    "\n",
    "class upsample(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, bilinear=False):\n",
    "        super(upsample, self).__init__()\n",
    "        \n",
    "        if bilinear:\n",
    "            self.up = nn.Upsample(scale_factor=2, mode='bilinear')\n",
    "        else:\n",
    "            self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        out = self.up(x)\n",
    "        return out\n",
    "    \n",
    "class concatenate_conv(nn.Module):\n",
    "    def __init__(self, layer_size):\n",
    "        super(concatenate_conv, self).__init__()\n",
    "        self.conv = double_conv_relu(layer_size*2, layer_size)\n",
    "        \n",
    "    def forward(self, encoder_layer, decoder_layer):\n",
    "        out = torch.cat([encoder_layer, decoder_layer], dim=1)\n",
    "        out = self.conv(out)\n",
    "        return out\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class unet(nn.Module):\n",
    "    def __init__(self, in_channels, out_classes, dropout=False):\n",
    "        super(unet, self).__init__()\n",
    "        \n",
    "        self.encoder_conv1 = double_conv_relu(in_channels, 64, dropout)\n",
    "        self.encoder_conv2 = double_conv_relu(64, 128, dropout)\n",
    "        self.encoder_conv3 = double_conv_relu(128, 256, dropout)\n",
    "        self.encoder_conv4 = double_conv_relu(256, 512, dropout)\n",
    "        self.encoder_conv5 = double_conv_relu(512, 1024, dropout) #set out channels to 512 instead of 1024 for memory\n",
    "        \n",
    "        self.decoder_conv1 = concatenate_conv(512)\n",
    "        self.decoder_conv2 = concatenate_conv(256)\n",
    "        self.decoder_conv3 = concatenate_conv(128)\n",
    "        self.decoder_conv4 = concatenate_conv(64)\n",
    "        \n",
    "        self.up1 = upsample(1024, 512)\n",
    "        self.up2 = upsample(512, 256)\n",
    "        self.up3 = upsample(256, 128)\n",
    "        self.up4 = upsample(128, 64)\n",
    "        \n",
    "        self.down = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        \n",
    "        self.output_conv = nn.Conv2d(64, out_classes, kernel_size=1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        encode1 = self.encoder_conv1(x)\n",
    "        out = self.down(encode1)\n",
    "        encode2 = self.encoder_conv2(out)\n",
    "        out = self.down(encode2)\n",
    "        encode3 = self.encoder_conv3(out)\n",
    "        out = self.down(encode3)\n",
    "        encode4 = self.encoder_conv4(out)\n",
    "        out = self.down(encode4)\n",
    "        encode5 = self.encoder_conv5(out)\n",
    "        decode = self.up1(encode5)\n",
    "        decode = self.decoder_conv1(encode4, decode)\n",
    "        decode = self.up2(decode)\n",
    "        decode = self.decoder_conv2(encode3, decode)\n",
    "        decode = self.up3(decode)\n",
    "        decode = self.decoder_conv3(encode2, decode)\n",
    "        decode = self.up4(decode)\n",
    "        decode = self.decoder_conv4(encode1, decode)\n",
    "        out = self.output_conv(decode)\n",
    "        \n",
    "        return out\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "31042434"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = unet(1,2)\n",
    "sum(p.numel() for p in model.parameters() if p.requires_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "def train_model(model, batch_size, epochs, lr=0.1, gpu=False):\n",
    "    \n",
    "    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    for epoch in range(epochs):\n",
    "        \n",
    "#         data_dir = os.path.join((os.getcwd()), 'data')\n",
    "#         labels = io.imread(os.path.join(data_dir, 'train-labels.tif')) #load training labels\n",
    "#         labels = ToTensor()(labels)\n",
    "#         labels.requires_grad = False\n",
    "#         labels = labels.transpose(0,1) #needed because of the TIF files\n",
    "        \n",
    "# #         labels = labels.unsqueeze(1)\n",
    "#         labels = labels[0]\n",
    "#         labels = labels.unsqueeze(0)\n",
    "#         labels = torch.Tensor.long(labels)\n",
    "#         labels = Variable(labels)\n",
    "        \n",
    "#         imgs = io.imread(os.path.join(data_dir, 'train-volume.tif')) #load training data\n",
    "#         imgs = ToTensor()(imgs)\n",
    "#         imgs = imgs.transpose(0,1)\n",
    "#         imgs.requires_grad = False\n",
    "#         imgs = imgs.unsqueeze(1)\n",
    "#         imgs = imgs[0]\n",
    "#         imgs = imgs.unsqueeze(0)\n",
    "#         imgs = Variable(imgs)\n",
    "#         if gpu:\n",
    "#             imgs = imgs.cuda()\n",
    "#             labels = labels.cuda()\n",
    "\n",
    "        x = Variable(torch.FloatTensor(np.random.random((2, 1, 256, 256))))\n",
    "            \n",
    "        \n",
    "        pred_masks = model(x)\n",
    "#         loss = criterion(pred_masks, labels)\n",
    "        loss = torch.sum(pred_masks)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        gc.collect()\n",
    "        del x, pred_masks\n",
    "        \n",
    "        epoch_loss += loss\n",
    "        print('Epoch {}, loss: {}'.format(epoch, epoch_loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, loss: Variable containing:\n",
      "-9346.3994\n",
      "[torch.FloatTensor of size 1]\n",
      "\n",
      "Epoch 1, loss: Variable containing:\n",
      "-1.5957e+13\n",
      "[torch.FloatTensor of size 1]\n",
      "\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-10-f70956ae98ac>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[0mmodel\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0munet\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m2\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m \u001b[0mtrain_model\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m5\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mgpu\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mFalse\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32m<ipython-input-9-1f70b10886b7>\u001b[0m in \u001b[0;36mtrain_model\u001b[1;34m(model, batch_size, epochs, lr, gpu)\u001b[0m\n\u001b[0;32m     37\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     38\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 39\u001b[1;33m         \u001b[0mpred_masks\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     40\u001b[0m \u001b[1;31m#         loss = criterion(pred_masks, labels)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     41\u001b[0m         \u001b[0mloss\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msum\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mpred_masks\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mD:\\Anaconda\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    355\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    356\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 357\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    358\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    359\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-6-a8137627abda>\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, x)\u001b[0m\n\u001b[0;32m     40\u001b[0m         \u001b[0mdecode\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdecoder_conv3\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mencode2\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdecode\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     41\u001b[0m         \u001b[0mdecode\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mup4\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdecode\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 42\u001b[1;33m         \u001b[0mdecode\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdecoder_conv4\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mencode1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdecode\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     43\u001b[0m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0moutput_conv\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdecode\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     44\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mD:\\Anaconda\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    355\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    356\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 357\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    358\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    359\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-3-faf05920d6f4>\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, encoder_layer, decoder_layer)\u001b[0m\n\u001b[0;32m     44\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mencoder_layer\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdecoder_layer\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     45\u001b[0m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mencoder_layer\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdecoder_layer\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 46\u001b[1;33m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mconv\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     47\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0mout\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     48\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mD:\\Anaconda\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    355\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    356\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 357\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    358\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    359\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-3-faf05920d6f4>\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, x)\u001b[0m\n\u001b[0;32m     12\u001b[0m         \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdropout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mdropout\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     13\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 14\u001b[1;33m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mconv1\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     15\u001b[0m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mnorm1\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     16\u001b[0m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mReLU\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mD:\\Anaconda\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    355\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    356\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 357\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    358\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    359\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mD:\\Anaconda\\lib\\site-packages\\torch\\nn\\modules\\conv.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, input)\u001b[0m\n\u001b[0;32m    280\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    281\u001b[0m         return F.conv2d(input, self.weight, self.bias, self.stride,\n\u001b[1;32m--> 282\u001b[1;33m                         self.padding, self.dilation, self.groups)\n\u001b[0m\u001b[0;32m    283\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    284\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mD:\\Anaconda\\lib\\site-packages\\torch\\nn\\functional.py\u001b[0m in \u001b[0;36mconv2d\u001b[1;34m(input, weight, bias, stride, padding, dilation, groups)\u001b[0m\n\u001b[0;32m     88\u001b[0m                 \u001b[0m_pair\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mgroups\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbackends\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcudnn\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbenchmark\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     89\u001b[0m                 torch.backends.cudnn.deterministic, torch.backends.cudnn.enabled)\n\u001b[1;32m---> 90\u001b[1;33m     \u001b[1;32mreturn\u001b[0m \u001b[0mf\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mweight\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     91\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     92\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "model = unet(1, 2)\n",
    "train_model(model, 1, 5, gpu=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = UNet(3, depth=5, merge_mode='concat', in_channels=1)\n",
    "x = Variable(torch.FloatTensor(np.random.random((1, 1, 512, 512))))\n",
    "out = model(x)\n",
    "loss = torch.sum(out)\n",
    "loss.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "del x, out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "del loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = os.path.join((os.getcwd()), 'data')\n",
    "labels = io.imread(os.path.join(data_dir, 'train-labels.tif')) #load training labels\n",
    "labels = torchvision.transforms.ToTensor()(labels)\n",
    "labels.requires_grad = False\n",
    "labels = labels.transpose(0,1) #needed because of the TIF files\n",
    "\n",
    "labels = torch.Tensor.long(labels)\n",
    "# labels = Variable(labels)\n",
    "\n",
    "imgs = io.imread(os.path.join(data_dir, 'train-volume.tif')) #load training data\n",
    "imgs = torchvision.transforms.ToTensor()(imgs)\n",
    "imgs = imgs.transpose(0,1)\n",
    "imgs.requires_grad = False\n",
    "imgs = imgs.unsqueeze(1)\n",
    "# imgs = Variable(imgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variable containing:\n",
      "( 0 ,.,.) = \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "     ...       ⋱       ...    \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "\n",
      "( 1 ,.,.) = \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "     ...       ⋱       ...    \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "\n",
      "( 2 ,.,.) = \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "     ...       ⋱       ...    \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "... \n",
      "\n",
      "(27 ,.,.) = \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "     ...       ⋱       ...    \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "\n",
      "(28 ,.,.) = \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "     ...       ⋱       ...    \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "\n",
      "(29 ,.,.) = \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "     ...       ⋱       ...    \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "[torch.LongTensor of size 30x512x512]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "torchvision.utils.save_image(labels[25], 'label.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = imgs[25]\n",
    "test = test.unsqueeze(0)\n",
    "test = Variable(test)\n",
    "out = model(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, indices = torch.max(out, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = np.array([[255/255,255/255,255/255], [100/255,100/255,100/255]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "colorsTensor = torch.Tensor(colors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = colorsTensor[indices.data.view(-1)].view(512, 512, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = output.permute(2, 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = torch.utils.data.TensorDataset(imgs.data, labels.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = torch.utils.data.DataLoader(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "( 0 ,.,.) = \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "     ...       ⋱       ...    \n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "   1   1   1  ...    1   1   1\n",
      "[torch.LongTensor of size 1x512x512]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for (batch, label) in dataloader:\n",
    "    print(label)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.unet import unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = unet(1,2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.6572)\n",
      "tensor(0.5146)\n",
      "tensor(0.4106)\n",
      "tensor(0.3620)\n",
      "tensor(0.3712)\n",
      "Epoch 0, loss: 2.315690279006958\n"
     ]
    }
   ],
   "source": [
    "train.train(model, imgs, labels, batch_size=1, epochs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'train' from 'D:\\\\Machine Learning\\\\UNet_pytorch\\\\train.py'>"
      ]
     },
     "execution_count": 172,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "importlib.reload(train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "( 0 , 0 ,.,.) = \n",
       "  0.4941  0.5412  0.5529  ...   0.6078  0.6392  0.5922\n",
       "  0.4196  0.4627  0.5294  ...   0.5843  0.6039  0.5529\n",
       "  0.4784  0.4824  0.5686  ...   0.6588  0.6510  0.6275\n",
       "           ...             ⋱             ...          \n",
       "  0.6314  0.6000  0.5647  ...   0.5451  0.6196  0.7176\n",
       "  0.6431  0.5843  0.5686  ...   0.4510  0.5529  0.6431\n",
       "  0.6941  0.6000  0.5373  ...   0.4706  0.5804  0.6706\n",
       "      ⋮  \n",
       "\n",
       "( 1 , 0 ,.,.) = \n",
       "  0.4510  0.4510  0.4627  ...   0.4902  0.5882  0.5451\n",
       "  0.5137  0.5176  0.4980  ...   0.5725  0.6627  0.5451\n",
       "  0.5098  0.5137  0.4314  ...   0.5647  0.6078  0.5020\n",
       "           ...             ⋱             ...          \n",
       "  0.4863  0.4118  0.4510  ...   0.6627  0.6941  0.7137\n",
       "  0.4902  0.4353  0.5137  ...   0.6745  0.6980  0.6706\n",
       "  0.5373  0.5608  0.5961  ...   0.6824  0.6745  0.6078\n",
       "      ⋮  \n",
       "\n",
       "( 2 , 0 ,.,.) = \n",
       "  0.6118  0.6431  0.6824  ...   0.6431  0.6431  0.6941\n",
       "  0.7255  0.7882  0.7686  ...   0.6392  0.6353  0.6863\n",
       "  0.7882  0.8078  0.8157  ...   0.6980  0.7333  0.7843\n",
       "           ...             ⋱             ...          \n",
       "  0.6471  0.5725  0.5765  ...   0.9725  0.9804  0.9882\n",
       "  0.6353  0.5529  0.6235  ...   0.9882  0.9961  0.9882\n",
       "  0.6275  0.5882  0.6275  ...   0.9647  0.9765  0.9843\n",
       "...     \n",
       "      ⋮  \n",
       "\n",
       "(27 , 0 ,.,.) = \n",
       "  0.3725  0.3961  0.4588  ...   0.6784  0.7216  0.5961\n",
       "  0.4118  0.5137  0.5922  ...   0.6118  0.7569  0.6745\n",
       "  0.2078  0.2431  0.3882  ...   0.4039  0.6902  0.6157\n",
       "           ...             ⋱             ...          \n",
       "  0.6980  0.6392  0.5843  ...   0.6157  0.5020  0.5725\n",
       "  0.6078  0.5373  0.5373  ...   0.6196  0.3490  0.3608\n",
       "  0.7294  0.6275  0.5922  ...   0.5176  0.2980  0.2863\n",
       "      ⋮  \n",
       "\n",
       "(28 , 0 ,.,.) = \n",
       "  0.2078  0.1843  0.1922  ...   0.6980  0.7020  0.6235\n",
       "  0.3961  0.4667  0.3961  ...   0.6784  0.6471  0.6902\n",
       "  0.7333  0.6667  0.5961  ...   0.6039  0.5804  0.8118\n",
       "           ...             ⋱             ...          \n",
       "  0.3725  0.3451  0.4667  ...   0.3137  0.3922  0.4588\n",
       "  0.3294  0.2510  0.2431  ...   0.5255  0.5490  0.5333\n",
       "  0.3922  0.2980  0.2863  ...   0.4235  0.6039  0.5843\n",
       "      ⋮  \n",
       "\n",
       "(29 , 0 ,.,.) = \n",
       "  0.2902  0.2471  0.1451  ...   0.5490  0.3686  0.3804\n",
       "  0.3529  0.2902  0.1686  ...   0.5294  0.4118  0.3529\n",
       "  0.2667  0.1882  0.1647  ...   0.4824  0.4863  0.3882\n",
       "           ...             ⋱             ...          \n",
       "  0.5451  0.4235  0.3961  ...   0.8353  0.7882  0.7804\n",
       "  0.4471  0.4196  0.3922  ...   0.7451  0.7608  0.7804\n",
       "  0.5176  0.4039  0.3137  ...   0.7255  0.7529  0.7882\n",
       "[torch.FloatTensor of size 30x1x512x512]"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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
}