{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load in \n\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport os\nimport keras\nimport time\nimport datetime\nfrom tensorflow.keras.applications.mobilenet import preprocess_input\nimport json\nimport cv2\nimport glob\n\nfrom torch.optim.optimizer import Optimizer, required\n\nfrom torch.autograd import Variable\nfrom torch import nn\nimport torch\nfrom torch import Tensor\nfrom torch.nn import Parameter\nimport tarfile\n\nimport torchvision      \nfrom torchvision.utils import save_image\nimport torchvision.transforms as transforms\n\nimport os\nprint(os.listdir(\"../input/\"))\n","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"df = pd.read_csv(\"../input/quickdraw-doodle-recognition/train_simplified/leaf.csv\")\ndf.to_csv(\"leaf.csv\")\nprint(len(df))\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 处理数据"},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef draw_cv2(raw_strokes, size=256, lw=6, time_color=True,BASE_SIZE = 256):\n    img = np.zeros((BASE_SIZE, BASE_SIZE), np.uint8)\n    for t, stroke in enumerate(raw_strokes):\n        for i in range(len(stroke[0]) - 1):\n            color = 255 - min(t, 10) * 13 if time_color else 255\n            _ = cv2.line(img, (stroke[0][i], stroke[1][i]),\n                         (stroke[0][i + 1], stroke[1][i + 1]), color, lw)\n    if size != BASE_SIZE:\n        return cv2.resize(img, (size, size))\n    else:\n        return img\n\ndef image_generator_xd(size, batchsize = 64, lw=6, n = 0,\n                       time_color=True, NCATS = 340,DP_DIR = '../input/shuffle-csvs/'):\n    paths = glob.glob(DP_DIR + \"*.csv*\")\n    paths.sort()\n    paths = paths[n:n+1]\n    paths = ['../input/quickdraw-doodle-recognition/train_simplified/leaf.csv']\n    print(paths)\n    while True:\n        for filename in paths:\n            for df in pd.read_csv(filename, chunksize=batchsize,engine='python'):\n                df['drawing'] = df['drawing'].apply(json.loads)\n                x = np.zeros((len(df), size, size))\n                for i, raw_strokes in enumerate(df.drawing.values):\n                    x[i, :, :] = draw_cv2(raw_strokes, size=size, lw=lw,\n                                             time_color=time_color)\n                x = preprocess_input(x).astype(np.float32)\n#                 y = keras.utils.to_categorical(df.y, num_classes=NCATS)\n                y = 0\n                yield x, y\n   ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ntrain_datagen = image_generator_xd(size=64, batchsize=128, n = 221,\n                                  DP_DIR = '../input/quickdraw-doodle-recognition/train_simplified/')\nx, y = next(train_datagen)\n# display_img(20,x)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def display_img(n,img):\n    if isinstance(img,torch.Tensor):\n        img = [np.rollaxis(i.numpy(), 0, 3)  for i in img]\n    for i in range(n):\n        plt.subplot(2,n//2,i+1)\n        plt.imshow(img[i])\n        plt.axis('off')\ndisplay_img(20,x)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"# 建立模型"},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef l2normalize(v, eps=1e-12):\n    return v / (v.norm() + eps)\n\nclass SpectralNorm(nn.Module):\n    def __init__(self, module, name='weight', power_iterations=1):\n        super(SpectralNorm, self).__init__()\n        self.module = module\n        self.name = name\n        self.power_iterations = power_iterations\n        if not self._made_params():\n            self._make_params()\n\n    def _update_u_v(self):\n        u = getattr(self.module, self.name + \"_u\")\n        v = getattr(self.module, self.name + \"_v\")\n        w = getattr(self.module, self.name + \"_bar\")\n\n        height = w.data.shape[0]\n        for _ in range(self.power_iterations):\n            v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))\n            u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))\n\n        # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))\n        sigma = u.dot(w.view(height, -1).mv(v))\n        setattr(self.module, self.name, w / sigma.expand_as(w))\n\n    def _made_params(self):\n        try:\n            u = getattr(self.module, self.name + \"_u\")\n            v = getattr(self.module, self.name + \"_v\")\n            w = getattr(self.module, self.name + \"_bar\")\n            return True\n        except AttributeError:\n            return False\n\n\n    def _make_params(self):\n        w = getattr(self.module, self.name)\n\n        height = w.data.shape[0]\n        width = w.view(height, -1).data.shape[1]\n\n        u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)\n        v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)\n        u.data = l2normalize(u.data)\n        v.data = l2normalize(v.data)\n        w_bar = Parameter(w.data)\n\n        del self.module._parameters[self.name]\n\n        self.module.register_parameter(self.name + \"_u\", u)\n        self.module.register_parameter(self.name + \"_v\", v)\n        self.module.register_parameter(self.name + \"_bar\", w_bar)\n\n\n    def forward(self, *args):\n        self._update_u_v()\n        return self.module.forward(*args)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n\nclass Self_Attn(nn.Module):\n    \"\"\" Self attention Layer\"\"\"\n    def __init__(self,in_dim,activation):\n        super(Self_Attn,self).__init__()\n        self.chanel_in = in_dim\n        self.activation = activation\n        \n        self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)\n        self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)\n        self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)\n        self.gamma = nn.Parameter(torch.zeros(1))\n\n        self.softmax  = nn.Softmax(dim=-1) #\n    def forward(self,x):\n        \"\"\"\n            inputs :\n                x : input feature maps( B X C X W X H)\n            returns :\n                out : self attention value + input feature \n                attention: B X N X N (N is Width*Height)\n        \"\"\"\n        m_batchsize,C,width ,height = x.size()\n        proj_query  = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)\n        proj_key =  self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)\n        energy =  torch.bmm(proj_query,proj_key) # transpose check\n        attention = self.softmax(energy) # BX (N) X (N) \n        proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N\n\n        out = torch.bmm(proj_value,attention.permute(0,2,1) )\n        out = out.view(m_batchsize,C,width,height)\n        \n        out = self.gamma*out + x\n        return out,attention\n\nclass Generator(nn.Module):\n    \"\"\"Generator.\"\"\"\n\n    def __init__(self, img_dim = 3, image_size=64, z_dim=100, conv_dim=64):\n        super(Generator, self).__init__()\n        self.imsize = image_size\n        layer1 = []\n        layer2 = []\n        layer3 = []\n        last = []\n\n        repeat_num = int(np.log2(self.imsize)) - 3\n        mult = 2 ** repeat_num # 8\n        layer1.append(SpectralNorm(nn.ConvTranspose2d(z_dim, conv_dim * mult, 4)))\n        layer1.append(nn.BatchNorm2d(conv_dim * mult))\n        layer1.append(nn.ReLU())\n\n        curr_dim = conv_dim * mult\n\n        layer2.append(SpectralNorm(nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))\n        layer2.append(nn.BatchNorm2d(int(curr_dim / 2)))\n        layer2.append(nn.ReLU())\n\n        curr_dim = int(curr_dim / 2)\n\n        layer3.append(SpectralNorm(nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))\n        layer3.append(nn.BatchNorm2d(int(curr_dim / 2)))\n        layer3.append(nn.ReLU())\n\n        if self.imsize == 64:\n            layer4 = []\n            curr_dim = int(curr_dim / 2)\n            layer4.append(SpectralNorm(nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))\n            layer4.append(nn.BatchNorm2d(int(curr_dim / 2)))\n            layer4.append(nn.ReLU())\n            self.l4 = nn.Sequential(*layer4)\n            curr_dim = int(curr_dim / 2)\n\n        self.l1 = nn.Sequential(*layer1)\n        self.l2 = nn.Sequential(*layer2)\n        self.l3 = nn.Sequential(*layer3)\n\n        last.append(nn.ConvTranspose2d(curr_dim, img_dim, 4, 2, 1))\n        last.append(nn.Tanh())\n        self.last = nn.Sequential(*last)\n\n        self.attn1 = Self_Attn( 128, 'relu')\n        self.attn2 = Self_Attn( 64,  'relu')\n\n    def forward(self, z):\n\n\n        \n        z = z.view(z.size(0), z.size(1), 1, 1) \n        # torch.Size([128, 128, 1, 1])\n        out=self.l1(z) # torch.Size([64, 512, 4, 4])\n        out=self.l2(out) # torch.Size([64, 256, 8, 8])\n        out=self.l3(out) # torch.Size([64, 128, 16, 16])\n        \n        out,p1 = self.attn1(out) # torch.Size([64, 128, 16, 16])  torch.Size([64, 256, 256])\n        \n        out=self.l4(out) # torch.Size([64, 64, 32, 32])  \n        \n        out,p2 = self.attn2(out) # torch.Size([64, 64, 32, 32])  torch.Size([64, 1024, 1024])\n        \n        out=self.last(out) # torch.Size([64, 3, 64, 64])\n\n        return out, p1, p2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Discriminator(nn.Module):\n    \"\"\"Discriminator, Auxiliary Classifier.\"\"\"\n\n    def __init__(self,  img_dim = 3,image_size=64, conv_dim=64):\n        super(Discriminator, self).__init__()\n        self.imsize = image_size\n        layer1 = []\n        layer2 = []\n        layer3 = []\n        last = []\n\n        layer1.append(SpectralNorm(nn.Conv2d(img_dim, conv_dim, 4, 2, 1)))\n        layer1.append(nn.LeakyReLU(0.1))\n\n        curr_dim = conv_dim\n\n        layer2.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))\n        layer2.append(nn.LeakyReLU(0.1))\n        curr_dim = curr_dim * 2\n\n        layer3.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))\n        layer3.append(nn.LeakyReLU(0.1))\n        curr_dim = curr_dim * 2\n\n        if self.imsize == 64:\n            layer4 = []\n            layer4.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))\n            layer4.append(nn.LeakyReLU(0.1))\n            self.l4 = nn.Sequential(*layer4)\n            curr_dim = curr_dim*2\n        self.l1 = nn.Sequential(*layer1)\n        self.l2 = nn.Sequential(*layer2)\n        self.l3 = nn.Sequential(*layer3)\n\n        last.append(nn.Conv2d(curr_dim, 1, 4))\n        self.last = nn.Sequential(*last)\n\n        self.attn1 = Self_Attn(256, 'relu')\n        self.attn2 = Self_Attn(512, 'relu')\n\n    def forward(self, x):\n#         torch.Size([128, 1, 64, 64])\n        out = self.l1(x) # torch.Size([200, 64, 32, 32])\n        out = self.l2(out) # torch.Size([200, 128, 16, 16])\n        out = self.l3(out) # torch.Size([200, 256, 8, 8])\n        \n        out,p1 = self.attn1(out) # torch.Size([200, 256, 8, 8]) torch.Size([200, 64, 64])\n        out=self.l4(out)# torch.Size([200, 512,4, 4])\n        out,p2 = self.attn2(out) # torch.Size([200, 512, 4, 4]) torch.Size([200, 16, 16])\n        \n        out=self.last(out) # torch.Size([200, 1, 1, 1])\n        \n\n        return out.squeeze(), p1, p2","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 训练部分"},{"metadata":{"trusted":true},"cell_type":"code","source":"adv_loss='wgan-gp'\nattn_path='./attn'\nbatch_size=64\nbeta1=0.0\nbeta2=0.9\nd_conv_dim=64\nd_iters=5\nd_lr=0.0004\ndataset='cifar'\ng_conv_dim=64 \ng_lr=0.0001\ng_num=5\nimage_path='./data'\nmodel_save_path='./models'\nlog_path='./logs'\nsample_path='./samples'\nimsize=64\nlambda_gp=10\nlr_decay=0.95\nmodel='sagan'\nnum_workers=2\nparallel=False\npretrained_model=None\nis_Cifar = True\ntotal_step=2000\nlog_step=total_step//100\nsample_step=total_step//100\nmodel_save_step=total_step//50\n\ntrain=True\nuse_tensorboard=False\nversion='sagan_1'\nz_dim=128\n\nfor i in [image_path,log_path,model_save_path,sample_path]:\n    if not os.path.exists(i):\n        os.mkdir(i)\n\ndef denorm(x):\n    out = (x + 1) / 2\n    return out.clamp_(0, 1)\n\ndef tensor2var(x, grad=False):\n    if not isinstance(x,torch.Tensor):\n        x = torch.Tensor(x).unsqueeze(1)\n    if torch.cuda.is_available():\n        x = x.cuda()\n    return Variable(x, requires_grad=grad)\n!ls","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# G = Generator(imsize, z_dim, g_conv_dim).cuda()\n# D = Discriminator(imsize, d_conv_dim).cuda()\n\n# def load_modle(G,D,path = None):\n#     if path is None:\n#         modle_load_path = glob.glob('../input/sagan/models/*.pth')\n#     modle_load_path.sort()\n#     if len(modle_load_path)>2:\n\n#         modle_load_path = modle_load_path[-2:]\n#         print(modle_load_path)\n#         D.load_state_dict(torch.load(modle_load_path[0]))\n#         G.load_state_dict(torch.load(modle_load_path[1]))\n#     else:\n#         print('file not found')\n#     return G,D\n# G,D = load_modle(G,D,path = None)\n\n# z = tensor2var(torch.randn(real_images.size(0), z_dim))\n\n# fake_images,gf1,gf2 = G(z)\n# d_out_fake,df1,df2 = D(fake_images)\n\n# # z.shape, fake_images.shape, d_out_fake.shape\n# # (torch.Size([64, 128]), torch.Size([64, 3, 64, 64]), torch.Size([64]))\n# display_img(10,fake_images.detach().cpu())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"模型生成与优化器选择"},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"\nis_Cifar = False\nif is_Cifar:\n    img_dim = 3\n    dataloader = trainloader\nelse:\n    img_dim = 1\n    dataloader = train_datagen\n    \nG = Generator(img_dim = img_dim ,image_size=imsize, z_dim = z_dim, conv_dim=g_conv_dim).cuda()\n\nD = Discriminator(img_dim = img_dim,image_size=imsize, conv_dim = d_conv_dim).cuda()\ntry:\n    G,D = load_modle(G,D,path = None)\nexcept:\n    print(\"load error\")\n\ng_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, G.parameters()), g_lr, [beta1, beta2])\nd_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, D.parameters()), d_lr, [beta1, beta2])\n\nc_loss = torch.nn.CrossEntropyLoss()\n\n# print(G)\n# print(D)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# 设置训练次数\n\n    \ndata_iter = iter(dataloader)\n# Start with trained model\nif pretrained_model:\n    start = pretrained_model + 1\nelse:\n    start = 0\n\n# Fixed input for debugging \nfixed_z = tensor2var(torch.randn(batch_size, z_dim))\nmodel_save_step","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"# Start time\nhist = []\nstart_time = time.time()\nfor step in range(start, total_step):\n\n    # ================== Train D ================== #\n    D.train()\n    G.train()\n\n    try:\n        real_images, _ = next(data_iter)\n    except:\n        print(\"error\")\n        data_iter = iter(dataloader)\n        real_images, _ = next(data_iter)\n\n    # Compute loss with real images\n    # dr1, dr2, df1, df2, gf1, gf2 are attention scores\n    real_images = tensor2var(real_images)\n    d_out_real,dr1,dr2 = D(real_images)\n    if adv_loss == 'wgan-gp':\n        d_loss_real = - torch.mean(d_out_real)\n    elif adv_loss == 'hinge':\n        d_loss_real = torch.nn.ReLU()(1.0 - d_out_real).mean()\n\n    # apply Gumbel Softmax\n    z = tensor2var(torch.randn(real_images.size(0), z_dim))\n    fake_images,gf1,gf2 = G(z)\n    d_out_fake,df1,df2 = D(fake_images)\n\n    if adv_loss == 'wgan-gp':\n        d_loss_fake = d_out_fake.mean()\n    elif adv_loss == 'hinge':\n        d_loss_fake = torch.nn.ReLU()(1.0 + d_out_fake).mean()\n\n\n    # Backward + Optimize\n    d_loss = d_loss_real + d_loss_fake\n    d_optimizer.zero_grad()\n    g_optimizer.zero_grad()\n    d_loss.backward()\n    d_optimizer.step()\n    \n\n    if adv_loss == 'wgan-gp':\n        # Compute gradient penalty \n        alpha = torch.rand(real_images.size(0), 1, 1, 1).cuda().expand_as(real_images)\n        interpolated = Variable(alpha * real_images.data + (1 - alpha) * fake_images.data,\n                                requires_grad=True)\n        out,_,_ = D(interpolated)\n\n        grad = torch.autograd.grad(outputs=out,\n                                   inputs=interpolated,\n                                   grad_outputs=torch.ones(out.size()).cuda(),\n                                   retain_graph=True,\n                                   create_graph=True,\n                                   only_inputs=True)[0]\n\n        grad = grad.view(grad.size(0), -1)\n        grad_l2norm = torch.sqrt(torch.sum(grad ** 2, dim=1))\n        d_loss_gp = torch.mean((grad_l2norm - 1) ** 2)\n\n        # Backward + Optimize\n        d_loss = lambda_gp * d_loss_gp\n\n        d_optimizer.zero_grad()\n#         g_optimizer.zero_grad()\n        d_loss.backward()\n        d_optimizer.step()\n\n    # ================== Train G and gumbel ================== #\n    # Create random noise\n    z = tensor2var(torch.randn(real_images.size(0), z_dim))\n    fake_images,_,_ = G(z)\n\n    # Compute loss with fake images\n    g_out_fake,_,_ = D(fake_images)  # batch x n\n    if adv_loss == 'wgan-gp':\n        g_loss_fake = - g_out_fake.mean()\n    elif adv_loss == 'hinge':\n        g_loss_fake = - g_out_fake.mean()\n\n#     d_optimizer.zero_grad()\n    g_optimizer.zero_grad()\n    g_loss_fake.backward()\n    g_optimizer.step()\n\n\n    # Print out log info\n    if (step + 1) % log_step == 0:\n        elapsed = time.time() - start_time\n        elapsed = str(datetime.timedelta(seconds=elapsed))\n        print(\"Elapsed [{}], G_step [{}/{}], D_step[{}/{}], d_out_real: {:.4f}, \"\n              \" ave_gamma_l3: {:.4f}, ave_gamma_l4: {:.4f}\".\n              format(elapsed, step + 1, total_step, (step + 1),\n                     total_step , d_loss_real.data,\n                     G.attn1.gamma.mean().data, G.attn2.gamma.mean().data ))\n\n    # Sample images\n    if (step + 1) % sample_step == 0:\n        print(str(step+1) + \"  |  Sample Saved..\")\n        fake_images,_,_= G(fixed_z)\n        save_image(denorm(fake_images.data),\n                   os.path.join(sample_path, '{:04}_fake.png'.format(step + 1)))\n\n    if (step+1) % model_save_step==0:\n        print(str(step+1) + \"  |  model Saved..\")\n        torch.save(G.state_dict(),\n                   os.path.join(model_save_path, '{:05}_G.pth'.format(step + 1)))\n        torch.save(D.state_dict(),\n                   os.path.join(model_save_path, '{:05}_D.pth'.format(step + 1)))\n   \n\n    hist.append([step,float(d_loss.data.cpu().numpy()),float(g_loss_fake.data.cpu().numpy())])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 绘制learning curve"},{"metadata":{"trusted":true},"cell_type":"code","source":"hist_ = np.array(hist)\nfig, axs = plt.subplots(2,1, figsize=(13,7))\n\nplt.subplot(2,1,1)\nplt.plot(hist_[:,1])\nplt.xlabel(\"step\")\nplt.ylabel(\"D loss\")\n\nplt.subplot(2,1,2)\nplt.plot(hist_[:,2])\nplt.xlabel(\"step\")\nplt.ylabel(\"G loss\")\n\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"im_paths = os.listdir(\"./samples/\")\n\nidx = [int(i[:-9]) for i in im_paths]\nidx.sort(reverse=True)\nbest_idx = idx[:6]\n\nbest_idx,len(im_paths)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, axs = plt.subplots(2,len(best_idx)//2, figsize=(13,9))\nfor i,idx in enumerate(best_idx):\n    plt.subplot(2,len(best_idx)//2,i+1)\n    a = plt.imread(\"./samples/\" + \"{:04}\".format(idx) + \"_fake.png\")\n\n    plt.imshow(a)\n    plt.axis('off')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}