27jun (1).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision import transforms, datasets\n",
    "import argparse\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "# torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
    "\n",
    "from torchvision.transforms import Compose, ToTensor, Resize\n",
    "import torch.utils.data as data\n",
    "from os import listdir\n",
    "from os.path import join\n",
    "from PIL import Image\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 32, 1, stride=1, padding=0)\n",
    "        for j in range(3):\n",
    "            for i in range(6):\n",
    "                exec(\"self.Resconv1\"+str(j)+str(i)+\"=\"+\"nn.Conv2d(32, 32, 3, stride=1, padding=1)\")\n",
    "                exec(\"self.Batch_norm1\"+str(j)+str(i)+\"=\"+\"nn.BatchNorm2d(32, track_running_stats=False)\")\n",
    "                exec(\"self.Resconv2\"+str(j)+str(i)+\"=\"+\"nn.Conv2d(32, 32, 3, stride=1, padding=1)\")\n",
    "                exec(\"self.Batch_norm2\"+str(j)+str(i)+\"=\"+\"nn.BatchNorm2d(32, track_running_stats=False)\")\n",
    "        self.Transconv1 = nn.ConvTranspose2d(32, 32, 3, stride=2, padding=1, output_padding = 1)\n",
    "        self.Transconv2 = nn.ConvTranspose2d(32, 32, 3, stride=2, padding=1, output_padding = 1)\n",
    "        self.conv2 = nn.Conv2d(32, 3*256, 1, stride=1, padding=0)\n",
    "\n",
    "        self.Maskconv1 = nn.Conv2d(3, 64, 7, stride=1, padding=3)\n",
    "        for i in range(20):\n",
    "            exec(\"self.Gateconv1\"+str(i)+ \"=\"+ \"nn.Conv2d(64, 2*64, 5, stride=1, padding=2)\")\n",
    "            exec(\"self.Gateconv2\"+str(i)+ \"=\"+ \"nn.Conv2d(2*64, 2*64, (1,1), stride=1, padding=0)\")\n",
    "            exec(\"self.Gateconv3\"+str(i)+ \"=\"+ \"nn.Conv2d(64, 2*64, (1,5), stride=1, padding=(0,2))\")\n",
    "            exec(\"self.Gateconv4\"+str(i)+ \"=\"+ \"nn.Conv2d(64, 64, (1,1), stride=1, padding=0)\")\n",
    "        self.Maskconv2 = nn.Conv2d(64, 1024, 1, stride=1, padding=0)\n",
    "        self.Maskconv3 = nn.Conv2d(1024, 3*256, 1, stride=1, padding=0)\n",
    "\n",
    "#         self.conv1.weight = nn.Parameter(self.weight_init(3, 32, (1,1), mask_type = None))\n",
    "#         self.Resconv1.weight = nn.Parameter(self.weight_init(32, 32, (3,3), mask_type = None))\n",
    "#         self.Transconv1.weight = nn.Parameter(self.weight_init(32, 32, (3,3), mask_type = None))\n",
    "#         self.Resconv2.weight = nn.Parameter(self.weight_init(32, 32, (3,3), mask_type = None))\n",
    "#         self.Transconv2.weight = nn.Parameter(self.weight_init(32, 32, (3,3), mask_type = None))\n",
    "#         self.conv2.weight = nn.Parameter(self.weight_init(32, 3*256, (1,1), mask_type = None))\n",
    "\n",
    "        self.Maskconv1.weight = nn.Parameter(self.weight_init(3, 64, (7,7), mask_type='A'))\n",
    "        for i in range(20):\n",
    "            exec(\"self.Gateconv1\"+str(i)+\".weight\"+\"=\"+\"nn.Parameter(self.weight_init(64, 2*64, (5,5), mask_type='C'))\")\n",
    "            exec(\"self.Gateconv2\"+str(i)+\".weight\"+\"=\"+\"nn.Parameter(self.weight_init(2*64, 2*64, (1,1), mask_type=None))\")\n",
    "            exec(\"self.Gateconv3\"+str(i)+\".weight\"+\"=\"+ \"nn.Parameter(self.weight_init(64, 2*64, (1,5), mask_type='B'))\")\n",
    "            exec(\"self.Gateconv4\"+str(i)+\".weight\"+\"=\"+\"nn.Parameter(self.weight_init(64, 64, (1,1), mask_type='B'))\")\n",
    "        self.Maskconv2.weight = nn.Parameter(self.weight_init(64, 1024, (1,1), mask_type='B'))\n",
    "        self.Maskconv3.weight = nn.Parameter(self.weight_init(1024, 3*256, (1,1), mask_type='B'))\n",
    "\n",
    "    def weight_init(self, in_channel, num_outputs, kernel_shape, mask_type=None):\n",
    "        kernel_h, kernel_w = kernel_shape\n",
    "        center_h = kernel_h // 2\n",
    "        center_w = kernel_w // 2\n",
    "\n",
    "        mask = np.zeros((num_outputs, in_channel, kernel_h, kernel_w), dtype=np.float32)\n",
    "        if mask_type is not None:\n",
    "            mask[:, :, :center_h, :] = 1\n",
    "            if mask_type == 'A':\n",
    "                mask[:, :, center_h, :center_w] = 1\n",
    "            if mask_type == 'B':\n",
    "                mask[:, :, center_h, :center_w+1] = 1\n",
    "        else:\n",
    "            mask[:, :, :, :] = 1\n",
    "\n",
    "        weights_shape = [num_outputs, in_channel, kernel_h, kernel_w]\n",
    "        weights = torch.empty(weights_shape, requires_grad=True, device=\"cuda\")\n",
    "        weights = nn.init.xavier_normal_(weights)\n",
    "        weights = weights*torch.from_numpy(mask).to(\"cuda\")\n",
    "        return weights\n",
    "\n",
    "    def prior_network(self, hr_images):\n",
    "        conv1 = self.Maskconv1(hr_images)\n",
    "        inputs = conv1\n",
    "        state = conv1\n",
    "        for i in range(20):\n",
    "            inputs, state = self.gated_conv2d(inputs, state, [5, 5], i)\n",
    "        conv2 = self.Maskconv2(inputs)\n",
    "        conv2 = F.relu(conv2)\n",
    "        prior_logits = self.Maskconv3(conv2)\n",
    "        prior_logits = prior_logits.permute(0, 2, 3, 1)\n",
    "        prior_logits = torch.cat((prior_logits[:, :, :, 0::3], prior_logits[:, :, :, 1::3], prior_logits[:, :, :, 2::3]), 1)\n",
    "        prior_logits = torch.permute(0, 3, 1, 2)\n",
    "        return prior_logits\n",
    "\n",
    "    def conditioning_network(self, lr_images):\n",
    "        res_num = 6\n",
    "        inputs = lr_images\n",
    "        inputs = self.conv1(inputs)\n",
    "        for i in range(2):\n",
    "            for j in range(res_num):\n",
    "                inputs = self.resnet_block(inputs, i, j)\n",
    "            inputs = eval(\"self.Transconv\"+str(i+1))(inputs)\n",
    "            inputs = F.relu(inputs)\n",
    "        for i in range(res_num):\n",
    "            inputs = self.resnet_block(inputs, 2, i)\n",
    "        conditioning_logits = self.conv2(inputs)\n",
    "        return conditioning_logits\n",
    "\n",
    "    def resnet_block(self, inputs, i, j):\n",
    "        conv1 = eval(\"self.Resconv1\"+str(i)+str(j))(inputs)\n",
    "        bn1 = eval(\"self.Batch_norm1\"+str(i)+str(j))(conv1)\n",
    "        relu1 = F.relu(bn1)\n",
    "        conv2 = eval(\"self.Resconv2\"+str(i)+str(j))(relu1)\n",
    "        bn2 = eval(\"self.Batch_norm2\"+str(i)+str(j))(conv2)\n",
    "        output = inputs + bn2\n",
    "        return output\n",
    "\n",
    "\n",
    "    def gated_conv2d(self, inputs, state, kernel_shape, i):\n",
    "        batch_size, in_channel, height, width  = list(inputs.size())\n",
    "        kernel_h, kernel_w = kernel_shape\n",
    "        left = eval(\"self.Gateconv1\"+str(i))(state)\n",
    "        left1 = left[:, 0:in_channel, :, :]\n",
    "        left2 = left[:, in_channel:, :, :]\n",
    "        left1 = F.tanh(left1)\n",
    "        left2 = F.sigmoid(left2)\n",
    "        new_state = left1 * left2\n",
    "        left2right = eval(\"self.Gateconv2\"+str(i))(left)\n",
    "        right = eval(\"self.Gateconv3\"+str(i))(inputs)\n",
    "        right = right + left2right\n",
    "        right1 = right[:, 0:in_channel, :, :]\n",
    "        right2 = right[:, in_channel:, :, :]\n",
    "        right1 = F.tanh(right1)\n",
    "        right2 = F.sigmoid(right2)\n",
    "        up_right = right1 * right2\n",
    "        up_right = eval(\"self.Gateconv4\"+str(i))(up_right)\n",
    "        outputs = inputs + up_right\n",
    "        return outputs, new_state\n",
    "\n",
    "    def forward(self, lr_images, hr_images):\n",
    "        hr_images = hr_images - 0.5\n",
    "        lr_images = lr_images - 0.5\n",
    "        prior_logits = self.prior_network(hr_images)\n",
    "        conditioning_logits = self.conditioning_network(lr_images)\n",
    "        return prior_logits, conditioning_logits\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def input_transform(crop_size):\n",
    "    return Compose([Resize((8,8),interpolation = 2),ToTensor(),])\n",
    "def target_transform(crop_size):\n",
    "    return Compose([Resize((32,32),interpolation = 2),ToTensor(),])\n",
    "def is_image_file(filename):\n",
    "    return any(filename.endswith(extension) for extension in [\".png\", \".jpg\", \".jpeg\"])\n",
    "def load_img(filepath):\n",
    "    img = Image.open(filepath)\n",
    "    return img\n",
    "\n",
    "class DatasetFromFolder(data.Dataset):\n",
    "    def __init__(self, image_dir, input_transform=None, target_transform=None):\n",
    "        super(DatasetFromFolder, self).__init__()\n",
    "        self.image_filenames = [join(image_dir, x) for x in listdir(image_dir) if is_image_file(x)]\n",
    "\n",
    "        self.input_transform = input_transform\n",
    "        self.target_transform = target_transform\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        inputs = load_img(self.image_filenames[index])\n",
    "        target = inputs.copy()\n",
    "        inputs = self.input_transform(inputs)\n",
    "        target = self.target_transform(target)\n",
    "\n",
    "        return inputs, target\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_filenames)\n",
    "    \n",
    "###\n",
    "def softmax_loss(logits, labels):\n",
    "    logits = logits.permute(0, 2, 3, 1)\n",
    "    logits = torch.reshape(logits, [-1, 256])\n",
    "    labels = labels.to(torch.int64)\n",
    "    labels = labels.permute(0, 2, 3, 1)\n",
    "    labels = torch.reshape(labels, [-1])\n",
    "    return F.cross_entropy(logits, labels)\n",
    "\n",
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader,1):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        prior_logits, conditioning_logits = model(lr_images = (data), hr_images = (target))\n",
    "        l1 = softmax_loss(prior_logits+conditioning_logits, torch.floor(target*255))\n",
    "        l2 = softmax_loss(conditioning_logits, torch.floor(target*255))\n",
    "        loss = l1+l2\n",
    "        loss3 = softmax_loss(prior_logits, torch.floor(target*255))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if batch_idx % 1 == 0:\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch, batch_idx * len(data), (len(train_loader)*len(data)),\n",
    "                100. * batch_idx / len(train_loader), loss.item()))\n",
    "\n",
    "        if batch_idx % 1000 == 0:\n",
    "            sample(model, data, target, len(data), mu=1.1, step=epoch*batch_idx)\n",
    "\n",
    "def test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (data, target) in enumerate(train_loader,1):\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            prior_logits, conditioning_logits = model(lr_images = (data), hr_images = (target))\n",
    "            l1 = softmax_loss(prior_logits+conditioning_logits, torch.floor(target*255))\n",
    "            l2 = softmax_loss(conditioning_logits, torch.floor(target*255))\n",
    "            test_loss += l1+l2 # sum up batch loss\n",
    "\n",
    "    test_loss /= len(test_loader)*32\n",
    "    print(\"test_loss : \", test_loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def logits_2_pixel_value(logits, mu=1.1):\n",
    "    rebalance_logits = logits * mu\n",
    "    probs = softmax(rebalance_logits)\n",
    "    pixel_dict = torch.arange(0, 256, dtype=torch.float32).to(\"cuda\")\n",
    "    pixels = torch.sum(probs*pixel_dict, dim=1)\n",
    "    return pixels/255\n",
    "\n",
    "def softmax(x):\n",
    "    \"\"\"Compute softmax values for each sets of scores in x.\"\"\"\n",
    "    a, b = torch.max(x, -1, keepdim=True, out=None)\n",
    "    e_x = torch.exp(x - a)\n",
    "    return e_x / e_x.sum(dim=-1, keepdim =True) # only difference\n",
    "\n",
    "def sample(model, data, target, batch_size, mu=1.1, step=None):\n",
    "    with torch.no_grad():\n",
    "        np_lr_imgs = data\n",
    "        np_hr_imgs = target\n",
    "        c_logits = model.conditioning_network\n",
    "        p_logits = model.prior_network\n",
    "        gen_hr_imgs = torch.zeros((batch_size, 3, 32, 32), dtype=torch.float32).to(\"cuda\")\n",
    "        np_c_logits = c_logits(torch.floor(np_lr_imgs*255))\n",
    "#         print(\"torch.floor(np_lr_imgs*255)\", torch.floor(np_lr_imgs*255)[0][:,0,0])\n",
    "#         print(\"c_logits\", np_c_logits[0].shape, torch.max(np_c_logits[0][:,0,0], -1))\n",
    "#         print(\"p_logits\", p_logits(gen_hr_imgs)[0].shape, p_logits(gen_hr_imgs)[0][:,0,0])\n",
    "#         print(\"target\", torch.floor(target*255)[0][:,0,0])\n",
    "        for i in range(32):\n",
    "            print(i)\n",
    "            for j in range(32):\n",
    "                for c in range(3):\n",
    "#                     print(i,j,c)\n",
    "                    np_p_logits = p_logits(gen_hr_imgs)\n",
    "                    new_pixel = logits_2_pixel_value(np_c_logits[:, c*256:(c+1)*256, i, j] + np_p_logits[:, c*256:(c+1)*256, i, j], mu=mu)\n",
    "                    gen_hr_imgs[:, c, i, j] = new_pixel\n",
    "        samples_dir = \"/home/eee/ug/15084005/DIH/samples/\"\n",
    "        print(\"sample\")\n",
    "        save_samples(np_lr_imgs, samples_dir + '/lr_' + str(mu*10) + '_' + str(step))\n",
    "        save_samples(np_hr_imgs, samples_dir + '/hr_' + str(mu*10) + '_' + str(step))\n",
    "        save_samples(gen_hr_imgs, samples_dir + '/generate_' + str(mu*10) + '_' + str(step))\n",
    "\n",
    "def save_samples(np_imgs, img_path):\n",
    "    print(\"save\")\n",
    "    torchvision.utils.save_image(np_imgs[0, :, :, :], img_path+\".png\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    use_gpu = True\n",
    "    num_epoch = 30\n",
    "    batch_size = 32\n",
    "    learning_rate = 4e-4\n",
    "\n",
    "    use_cuda = use_gpu and torch.cuda.is_available()\n",
    "    device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "    torch.cuda.set_device(0)\n",
    "\n",
    "    image_dir = '/home/eee/ug/15084005/DIH/CelebA/CelebA/train/img_align_celeba/'\n",
    "    train_set = DatasetFromFolder(image_dir,input_transform=input_transform(1),target_transform=target_transform(1))\n",
    "    train_loader = DataLoader(dataset=train_set, batch_size = batch_size, shuffle=True)\n",
    "    \n",
    "    image_dir2 = '/home/eee/ug/15084005/DIH/CelebA/CelebA/test/data/'\n",
    "    test_set = DatasetFromFolder(image_dir2,input_transform=input_transform(1),target_transform=target_transform(1))\n",
    "    test_loader = DataLoader(dataset=train_set, batch_size = batch_size, shuffle=True)\n",
    "\n",
    "    model = Net()\n",
    "    model = model.to(device)\n",
    "\n",
    "    optimizer = torch.optim.RMSprop(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.95)\n",
    "    print(sum(p.numel() for p in model.parameters() if p.requires_grad))\n",
    "    for epoch in range(1, num_epoch+1):\n",
    "        train(model, device, train_loader, optimizer, epoch)\n",
    "        # sample(model, train_loader_lr, train_loader_hr, mu=1.1, step=epoch)\n",
    "        test(model, device, test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "cuda runtime error (46) : all CUDA-capable devices are busy or unavailable at /opt/conda/conda-bld/pytorch_1524584710464/work/aten/src/THC/THCTensorRandom.cu:25",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-6-263240bbee7e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mmain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-5-6d737bd0653d>\u001b[0m in \u001b[0;36mmain\u001b[0;34m()\u001b[0m\n\u001b[1;32m     17\u001b[0m     \u001b[0mtest_loader\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mDataLoader\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrain_set\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshuffle\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 19\u001b[0;31m     \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mNet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     20\u001b[0m     \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-2-f0068fc2dfcf>\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;31m#         self.conv2.weight = nn.Parameter(self.weight_init(32, 3*256, (1,1), mask_type = None))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 31\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mMaskconv1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mParameter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight_init\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m64\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;36m7\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m7\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmask_type\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'A'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     32\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m20\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     33\u001b[0m             \u001b[0mexec\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"self.Gateconv1\"\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m\".weight\"\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m\"=\"\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m\"nn.Parameter(self.weight_init(64, 2*64, (5,5), mask_type='C'))\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-2-f0068fc2dfcf>\u001b[0m in \u001b[0;36mweight_init\u001b[0;34m(self, in_channel, num_outputs, kernel_shape, mask_type)\u001b[0m\n\u001b[1;32m     54\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     55\u001b[0m         \u001b[0mweights_shape\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mnum_outputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_channel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkernel_h\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkernel_w\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 56\u001b[0;31m         \u001b[0mweights\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mempty\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mweights_shape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrequires_grad\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"cuda\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     57\u001b[0m         \u001b[0mweights\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minit\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mxavier_normal_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mweights\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     58\u001b[0m         \u001b[0mweights\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mweights\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_numpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmask\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"cuda\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/cuda/__init__.py\u001b[0m in \u001b[0;36m_lazy_init\u001b[0;34m()\u001b[0m\n\u001b[1;32m    159\u001b[0m             \"Cannot re-initialize CUDA in forked subprocess. \" + msg)\n\u001b[1;32m    160\u001b[0m     \u001b[0m_check_driver\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 161\u001b[0;31m     \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_C\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_cuda_init\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    162\u001b[0m     \u001b[0m_cudart\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_load_cudart\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    163\u001b[0m     \u001b[0m_cudart\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcudaGetErrorName\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrestype\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mctypes\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mc_char_p\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: cuda runtime error (46) : all CUDA-capable devices are busy or unavailable at /opt/conda/conda-bld/pytorch_1524584710464/work/aten/src/THC/THCTensorRandom.cu:25"
     ]
    }
   ],
   "source": [
    "main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pytorch]",
   "language": "python",
   "name": "conda-env-pytorch-py"
  },
  "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
}