{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Import Libraries","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport matplotlib.pyplot as plt\nimport matplotlib.image as mpimg\n%matplotlib inline\nimport numpy as np\nimport torch\nfrom torch import nn, optim\nfrom torch import autograd\nimport torch.nn.functional as F\nfrom torch.nn import Parameter\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nfrom torch.utils.data import Dataset,DataLoader,Subset\n\nimport cv2\nimport albumentations as A\n\nimport glob\nimport shutil\nfrom tqdm.auto import tqdm\nimport time\nimport random\n\nimport os\nimport json\nimport math\nimport concurrent.futures","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-01-16T12:16:36.436236Z","iopub.execute_input":"2022-01-16T12:16:36.436594Z","iopub.status.idle":"2022-01-16T12:16:41.493661Z","shell.execute_reply.started":"2022-01-16T12:16:36.436506Z","shell.execute_reply":"2022-01-16T12:16:41.492598Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TIME_LIMIT = 32400 - 60*15\nstart_time = time.time()\ndef elapsed_time(start_time):\n    return time.time() - start_time","metadata":{"execution":{"iopub.status.busy":"2022-01-16T12:16:41.496506Z","iopub.execute_input":"2022-01-16T12:16:41.496847Z","iopub.status.idle":"2022-01-16T12:16:41.503319Z","shell.execute_reply.started":"2022-01-16T12:16:41.496801Z","shell.execute_reply":"2022-01-16T12:16:41.502252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#random seeds\nseed = 2021\nrandom.seed(seed)\ntorch.manual_seed(seed)\ntorch.cuda.manual_seed_all(seed)\nnp.random.seed(seed)\nos.environ['PYTHONHASHSEED'] = str(seed)\n\nBATCH_SIZE  = 32\nNUM_WORKERS = 4\nEMA = False\nLABEL_NOISE = False\nLABEL_NOISE_PROB = 0.1","metadata":{"execution":{"iopub.status.busy":"2022-01-16T12:16:41.504911Z","iopub.execute_input":"2022-01-16T12:16:41.506061Z","iopub.status.idle":"2022-01-16T12:16:41.519303Z","shell.execute_reply.started":"2022-01-16T12:16:41.506014Z","shell.execute_reply":"2022-01-16T12:16:41.518078Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"classes = [\"Eevee\", \"Espeon\", \"Flareon\", \"Glaceon\", \"Jolteon\", \"Leafeon\", \"Sylveon\", \"Umbreon\", \"Vaporeon\"]\nn_classes = classes.__len__()\nimg_size = 64","metadata":{"execution":{"iopub.status.busy":"2022-01-16T12:16:41.521494Z","iopub.execute_input":"2022-01-16T12:16:41.522540Z","iopub.status.idle":"2022-01-16T12:16:41.528402Z","shell.execute_reply.started":"2022-01-16T12:16:41.522491Z","shell.execute_reply":"2022-01-16T12:16:41.527237Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Pre-Processing","metadata":{}},{"cell_type":"code","source":"json_paths = []\nimage_paths = []\next = {}\nimg_extensions = ['.jpg', '.jpeg', '.PNG', '.JPG']\nfor root, dirs, files in os.walk(\"../input/evdataset/\"):\n    for name in files:\n        filename, extension = os.path.splitext(name)\n        ext[extension] = \"ok\"\n\n        if (\".json\" in extension) and (\"Metadata\" in root) and (\"Grouped\" in root):\n            json_paths.append(os.path.join(root, name))\n            \n        if (extension in img_extensions) and (\"Data\" in root) and (\"Grouped\" not in root) and (\"NotEevee\" not in root) and (\"Unsorted\" not in root):\n            image_paths.append(os.path.join(root, name))","metadata":{"execution":{"iopub.status.busy":"2022-01-16T12:16:41.532998Z","iopub.execute_input":"2022-01-16T12:16:41.533913Z","iopub.status.idle":"2022-01-16T12:16:45.519356Z","shell.execute_reply.started":"2022-01-16T12:16:41.533850Z","shell.execute_reply":"2022-01-16T12:16:45.518343Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def json_to_imagedata(json_path):\n    X = []\n    y = []\n    with open(json_path, 'r') as f:\n        data = json.load(f)\n#     if data[\"flags\"].get(\"monochrome\"):\n#         continue\n    img_name = os.path.split(data['imagePath'])[-1]\n    img_path = os.path.join(os.path.dirname(json_path), data['imagePath'].replace('\\\\', '/'))\n    img_path = os.path.normpath(img_path)\n    img = cv2.imread(img_path)\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n    # pad\n    pad_y = img.shape[0]\n    pad_x = img.shape[1]\n    padded = cv2.copyMakeBorder(img, pad_y, pad_y, pad_x, pad_x, cv2.BORDER_REFLECT)\n\n    for mugshot in data['shapes']:\n        if mugshot['shape_type'] != 'circle':\n            continue\n        if mugshot[\"flags\"].get(\"obscured\"):\n            continue\n        if mugshot[\"label\"] == \"Eevee\" and data[\"flags\"].get(\"monochrome\"):\n            continue\n\n        p1 = mugshot['points'][0] # center\n        p1[0] += pad_x\n        p1[1] += pad_y\n        p2 = mugshot['points'][1]\n        p2[0] += pad_x\n        p2[1] += pad_y\n\n        # rotato\n        adj_len = abs(p1[0] - p2[0])\n        oppo_len= abs(p1[1] - p2[1])\n        angle = math.degrees(math.atan(adj_len/oppo_len))\n        if p1[0] < p2[0]: # p1 is left of p2\n            if p1[1] < p2[1]: # p1 is top of p2\n                angle = -angle\n            else:\n                angle = angle+180\n        else: # p1 is right of p2\n            if p1[1] < p2[1]: # p1 is top of p2\n                angle = angle\n            else:\n                angle = -angle+180\n        rot_mat = cv2.getRotationMatrix2D(p1, angle, 1.0)\n        img_rotated = cv2.warpAffine(padded, rot_mat, padded.shape[1::-1], flags=cv2.INTER_LINEAR)\n\n        # crop\n        hypo_len = math.sqrt(adj_len**2 + oppo_len**2) *1.5# larger\n        top_left = (int(p1[0] - hypo_len), int(p1[1] - hypo_len))\n        sqr_len = int(hypo_len * 2)\n        img_crop = img_rotated[top_left[1]:top_left[1] + sqr_len, top_left[0]:top_left[0] + sqr_len]\n\n        # resize\n        img_resize = cv2.resize(img_crop, (img_size, img_size))\n        \n        X.append(img_resize)\n\n        ind = classes.index(mugshot[\"label\"])\n#         onehot = [0.0 for x in classes]\n#         onehot[ind] = 1.0\n        y.append(ind)\n    return X, y\ndef image_path_to_img_data(img_path):\n    img = cv2.imread(img_path)\n#     if img is None: return None, None\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    # crop largest inner square\n    center = np.array(img.shape) / 2\n    h, w, c = img.shape\n    x = center[1] - w/2\n    y = center[0] - h/2\n    img_crop = img[int(y):int(y+h), int(x):int(x+w)]\n    img_resize = cv2.resize(img_crop, (img_size, img_size))\n    \n    img_class = img_path.split(\"/\")[-2]\n    ind = classes.index(img_class)\n    return img_resize, ind","metadata":{"execution":{"iopub.status.busy":"2022-01-16T12:16:45.521344Z","iopub.execute_input":"2022-01-16T12:16:45.522045Z","iopub.status.idle":"2022-01-16T12:16:45.547432Z","shell.execute_reply.started":"2022-01-16T12:16:45.522001Z","shell.execute_reply":"2022-01-16T12:16:45.546470Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with concurrent.futures.ProcessPoolExecutor() as executor:\n    Xy_pairs = list(tqdm(\n        executor.map(image_path_to_img_data, image_paths),\n        total=image_paths.__len__()\n    ))\n    list_of_Xy_pairs_in_list = list(tqdm(\n        executor.map(json_to_imagedata, json_paths),\n        total=json_paths.__len__()\n    ))\nX = []\ny = []\nfor Xy_pairs_in_list in list_of_Xy_pairs_in_list:\n    X.extend(Xy_pairs_in_list[0])\n    y.extend(Xy_pairs_in_list[1])\nfor pair in Xy_pairs:\n    X.append(pair[0])\n    y.append(pair[1])","metadata":{"execution":{"iopub.status.busy":"2022-01-16T12:16:45.549802Z","iopub.execute_input":"2022-01-16T12:16:45.550634Z","iopub.status.idle":"2022-01-16T12:32:02.512410Z","shell.execute_reply.started":"2022-01-16T12:16:45.550588Z","shell.execute_reply":"2022-01-16T12:32:02.511362Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cnt = [y.count(item) for item in range(9)]\nvals = [item for item in range(9)]\nplt.bar(vals, cnt)\nplt.title(\"Class distribution\")\nplt.show()\ncnt","metadata":{"execution":{"iopub.status.busy":"2022-01-16T12:32:02.514585Z","iopub.execute_input":"2022-01-16T12:32:02.514903Z","iopub.status.idle":"2022-01-16T12:32:02.782065Z","shell.execute_reply.started":"2022-01-16T12:32:02.514860Z","shell.execute_reply":"2022-01-16T12:32:02.781100Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X = []\ny = []\nfor Xy_pairs_in_list in list_of_Xy_pairs_in_list:\n    X.extend(Xy_pairs_in_list[0])\n    y.extend(Xy_pairs_in_list[1])","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:39:36.992821Z","iopub.execute_input":"2022-01-16T14:39:36.993711Z","iopub.status.idle":"2022-01-16T14:39:37.003499Z","shell.execute_reply.started":"2022-01-16T14:39:36.993651Z","shell.execute_reply":"2022-01-16T14:39:37.002241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(X[0])","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:05.844454Z","iopub.execute_input":"2022-01-16T14:40:05.844865Z","iopub.status.idle":"2022-01-16T14:40:06.330611Z","shell.execute_reply.started":"2022-01-16T14:40:05.844815Z","shell.execute_reply":"2022-01-16T14:40:06.328158Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DogDataset(Dataset):\n    def __init__(self, X, y, transform1=None, transform2=None):\n        self.transform1 = transform1\n        self.transform2 = transform2\n        \n        self.imgs   = []\n        self.labels = y\n            \n        self.transform = transforms.ToTensor()\n        for img in X:\n            self.imgs.append(self.transform(img))\n            \n    def __len__(self):\n        return len(self.imgs)\n    \n    def __getitem__(self,idx):\n        img = self.imgs[idx]\n        if self.transform2:\n            img = self.transform2(img)\n        label = self.labels[idx]\n        return {'img':img, 'label':label}","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:06.333915Z","iopub.execute_input":"2022-01-16T14:40:06.334307Z","iopub.status.idle":"2022-01-16T14:40:06.342682Z","shell.execute_reply.started":"2022-01-16T14:40:06.334262Z","shell.execute_reply":"2022-01-16T14:40:06.341657Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n# generate 64x64 images!\n#resize_size = 68\nbatch_size  = BATCH_SIZE\nMEAN1,MEAN2,MEAN3 = 0.5, 0.5, 0.5\nSTD1,STD2,STD3    = 0.5, 0.5, 0.5\n\ntransform1 = transforms.Compose([transforms.Resize(img_size)])\n\ntransform2 = transforms.Compose([transforms.RandomCrop(img_size),\n                                 transforms.RandomAffine(degrees=5),\n                                 transforms.RandomHorizontalFlip(p=0.5),\n                                 #transforms.RandomApply(random_transforms, p=0.3),\n                                 transforms.Normalize(mean=[MEAN1, MEAN2, MEAN3],\n                                                      std=[STD1, STD2, STD3]),\n                                ])\n\ntrain_set = DogDataset(X, y,\n                       transform1=transform1,\n                       transform2=transform2,)\n\ntrain_loader = DataLoader(train_set,\n                          shuffle=True, batch_size=batch_size,\n                          num_workers=NUM_WORKERS, pin_memory=True)","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:06.344523Z","iopub.execute_input":"2022-01-16T14:40:06.345168Z","iopub.status.idle":"2022-01-16T14:40:07.017110Z","shell.execute_reply.started":"2022-01-16T14:40:06.345121Z","shell.execute_reply":"2022-01-16T14:40:07.015931Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(train_set)","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:07.019412Z","iopub.execute_input":"2022-01-16T14:40:07.019650Z","iopub.status.idle":"2022-01-16T14:40:07.029162Z","shell.execute_reply.started":"2022-01-16T14:40:07.019619Z","shell.execute_reply":"2022-01-16T14:40:07.027950Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Device","metadata":{}},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndevice","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:07.033543Z","iopub.execute_input":"2022-01-16T14:40:07.034055Z","iopub.status.idle":"2022-01-16T14:40:07.044824Z","shell.execute_reply.started":"2022-01-16T14:40:07.034023Z","shell.execute_reply":"2022-01-16T14:40:07.043608Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"def count_parameters(model):\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:07.046862Z","iopub.execute_input":"2022-01-16T14:40:07.047356Z","iopub.status.idle":"2022-01-16T14:40:07.054277Z","shell.execute_reply.started":"2022-01-16T14:40:07.047315Z","shell.execute_reply":"2022-01-16T14:40:07.052825Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def conv3x3(in_channel, out_channel): #not change resolusion\n    return nn.Conv2d(in_channel,out_channel,\n                      kernel_size=3,stride=1,padding=1,dilation=1,bias=False)\n\ndef conv1x1(in_channel, out_channel): #not change resolution\n    return nn.Conv2d(in_channel,out_channel,\n                      kernel_size=1,stride=1,padding=0,dilation=1,bias=False)\n\ndef init_weight(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        nn.init.orthogonal_(m.weight, gain=1)\n        if m.bias is not None:\n            m.bias.data.zero_()\n            \n    elif classname.find('Batch') != -1:\n        m.weight.data.normal_(1,0.02)\n        m.bias.data.zero_()\n    \n    elif classname.find('Linear') != -1:\n        nn.init.orthogonal_(m.weight, gain=1)\n        if m.bias is not None:\n            m.bias.data.zero_()\n    \n    elif classname.find('Embedding') != -1:\n        nn.init.orthogonal_(m.weight, gain=1)\n        \n\nclass Attention(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.channels = channels\n        self.theta    = nn.utils.spectral_norm(conv1x1(channels, channels//8)).apply(init_weight)\n        self.phi      = nn.utils.spectral_norm(conv1x1(channels, channels//8)).apply(init_weight)\n        self.g        = nn.utils.spectral_norm(conv1x1(channels, channels//2)).apply(init_weight)\n        self.o        = nn.utils.spectral_norm(conv1x1(channels//2, channels)).apply(init_weight)\n        self.gamma    = nn.Parameter(torch.tensor(0.), requires_grad=True)\n        \n    def forward(self, inputs):\n        batch,c,h,w = inputs.size()\n        theta = self.theta(inputs) #->(*,c/8,h,w)\n        phi   = F.max_pool2d(self.phi(inputs), [2,2]) #->(*,c/8,h/2,w/2)\n        g     = F.max_pool2d(self.g(inputs), [2,2]) #->(*,c/2,h/2,w/2)\n        \n        theta = theta.view(batch, self.channels//8, -1) #->(*,c/8,h*w)\n        phi   = phi.view(batch, self.channels//8, -1) #->(*,c/8,h*w/4)\n        g     = g.view(batch, self.channels//2, -1) #->(*,c/2,h*w/4)\n        \n        beta = F.softmax(torch.bmm(theta.transpose(1,2), phi), -1) #->(*,h*w,h*w/4)\n        o    = self.o(torch.bmm(g, beta.transpose(1,2)).view(batch,self.channels//2,h,w)) #->(*,c,h,w)\n        return self.gamma*o + inputs\n        \n    \nclass ConditionalNorm(nn.Module):\n    def __init__(self, in_channel, n_condition):\n        super().__init__()\n        self.bn = nn.BatchNorm2d(in_channel, affine=False) #no learning parameters\n        self.embed = nn.Linear(n_condition, in_channel* 2)\n        \n        nn.init.orthogonal_(self.embed.weight.data[:, :in_channel], gain=1)\n        self.embed.weight.data[:, in_channel:].zero_()\n\n    def forward(self, inputs, label):\n        out = self.bn(inputs)\n        embed = self.embed(label.float())\n        gamma, beta = embed.chunk(2, dim=1)\n        gamma = gamma.unsqueeze(2).unsqueeze(3)\n        beta = beta.unsqueeze(2).unsqueeze(3)\n        out = gamma * out + beta\n        return out","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:07.056627Z","iopub.execute_input":"2022-01-16T14:40:07.057451Z","iopub.status.idle":"2022-01-16T14:40:07.084566Z","shell.execute_reply.started":"2022-01-16T14:40:07.057407Z","shell.execute_reply":"2022-01-16T14:40:07.083324Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#BigGAN + leaky_relu           \nclass ResBlock_G(nn.Module):\n    def __init__(self, in_channel, out_channel, condition_dim, upsample=True):\n        super().__init__()\n        self.cbn1 = ConditionalNorm(in_channel, condition_dim)\n        self.upsample = nn.Sequential()\n        if upsample:\n            self.upsample.add_module('upsample',nn.Upsample(scale_factor=2, mode='nearest'))\n        self.conv3x3_1 = nn.utils.spectral_norm(conv3x3(in_channel, out_channel)).apply(init_weight)\n        self.cbn2 = ConditionalNorm(out_channel, condition_dim)\n        self.conv3x3_2 = nn.utils.spectral_norm(conv3x3(out_channel, out_channel)).apply(init_weight) \n        self.conv1x1   = nn.utils.spectral_norm(conv1x1(in_channel, out_channel)).apply(init_weight)\n        \n    def forward(self, inputs, condition):\n        x  = F.leaky_relu(self.cbn1(inputs, condition))\n        x  = self.upsample(x)\n        x  = self.conv3x3_1(x)\n        x  = self.conv3x3_2(F.leaky_relu(self.cbn2(x, condition)))\n        x += self.conv1x1(self.upsample(inputs)) #shortcut\n        return x\n\nclass Generator(nn.Module):\n    def __init__(self, n_feat, codes_dim=24, n_classes=n_classes):\n        super().__init__()\n        self.fc   = nn.Sequential(\n            nn.utils.spectral_norm(nn.Linear(codes_dim, 16*n_feat*4*4)).apply(init_weight)\n        )\n        self.res1 = ResBlock_G(16*n_feat, 16*n_feat, codes_dim+n_classes, upsample=True)\n        self.res2 = ResBlock_G(16*n_feat,  8*n_feat, codes_dim+n_classes, upsample=True)\n        #self.attn2 = Attention(8*n_feat)\n        self.res3 = ResBlock_G( 8*n_feat,  4*n_feat, codes_dim+n_classes, upsample=True)\n        self.attn = Attention(4*n_feat)\n        self.res4 = ResBlock_G( 4*n_feat,  2*n_feat, codes_dim+n_classes, upsample=True)\n        self.conv = nn.Sequential(\n            #nn.BatchNorm2d(2*n_feat).apply(init_weight),\n            nn.LeakyReLU(),\n            nn.utils.spectral_norm(conv3x3(2*n_feat, 3)).apply(init_weight),\n        )\n        \n    def forward(self, z, label_ohe, codes_dim=24):\n        '''\n        z.shape = (*,120)\n        label_ohe.shape = (*,n_classes)\n        '''\n        batch = z.size(0)\n        z = z.squeeze()\n        label_ohe = label_ohe.squeeze()\n        codes = torch.split(z, codes_dim, dim=1)\n        \n        x = self.fc(codes[0]) #->(*,16ch*4*4)\n        x = x.view(batch,-1,4,4) #->(*,16ch,4,4)\n        \n        condition = torch.cat([codes[1], label_ohe], dim=1) #(*,codes_dim+n_classes)\n        x = self.res1(x, condition)#->(*,16ch,8,8)\n        \n        condition = torch.cat([codes[2], label_ohe], dim=1)\n        x = self.res2(x, condition) #->(*,8ch,16,16)\n        #x = self.attn2(x) #not change shape\n        \n        condition = torch.cat([codes[3], label_ohe], dim=1)\n        x = self.res3(x, condition) #->(*,4ch,32,32)\n        x = self.attn(x) #not change shape\n        \n        condition = torch.cat([codes[4], label_ohe], dim=1)\n        x = self.res4(x, condition) #->(*,2ch,64,64)\n        \n        x = self.conv(x) #->(*,3,64,64)\n        x = torch.tanh(x)\n        return x\n    \n\nclass ResBlock_D(nn.Module):\n    def __init__(self, in_channel, out_channel, downsample=True):\n        super().__init__()\n        self.layer = nn.Sequential(\n            nn.LeakyReLU(0.2),\n            nn.utils.spectral_norm(conv3x3(in_channel, out_channel)).apply(init_weight),\n            nn.LeakyReLU(0.2),\n            nn.utils.spectral_norm(conv3x3(out_channel, out_channel)).apply(init_weight),\n        )\n        self.shortcut = nn.Sequential(\n            nn.utils.spectral_norm(conv1x1(in_channel,out_channel)).apply(init_weight),\n        )\n        if downsample:\n            self.layer.add_module('avgpool',nn.AvgPool2d(kernel_size=2,stride=2))\n            self.shortcut.add_module('avgpool',nn.AvgPool2d(kernel_size=2,stride=2))\n        \n    def forward(self, inputs):\n        x  = self.layer(inputs)\n        x += self.shortcut(inputs)\n        return x\n    \n\nclass Discriminator(nn.Module):\n    def __init__(self, n_feat, n_classes=n_classes):\n        super().__init__()\n        self.res1 = ResBlock_D(3, n_feat, downsample=True)\n        self.attn = Attention(n_feat)\n        self.res2 = ResBlock_D(  n_feat, 2*n_feat, downsample=True)\n        #self.attn2 = Attention(2*n_feat)\n        self.res3 = ResBlock_D(2*n_feat, 4*n_feat, downsample=True)\n        self.res4 = ResBlock_D(4*n_feat, 8*n_feat, downsample=True)\n        self.res5 = ResBlock_D(8*n_feat,16*n_feat, downsample=False)\n        self.fc   = nn.utils.spectral_norm(nn.Linear(16*n_feat,1)).apply(init_weight)\n        self.embedding = nn.Embedding(num_embeddings=n_classes, embedding_dim=16*n_feat).apply(init_weight)\n        \n    def forward(self, inputs, label):\n        batch = inputs.size(0) #(*,3,64,64)\n        h = self.res1(inputs) #->(*,ch,32,32)\n        h = self.attn(h) #not change shape\n        h = self.res2(h) #->(*,2ch,16,16)\n        #h = self.attn2(h) #not change shape\n        h = self.res3(h) #->(*,4ch,8,8)\n        h = self.res4(h) #->(*,8ch,4,4)\n        h = self.res5(h) #->(*,16ch,4,4)\n        h = torch.sum((F.leaky_relu(h,0.2)).view(batch,-1,4*4), dim=2) #GlobalSumPool ->(*,16ch)\n        outputs = self.fc(h) #->(*,1)\n        \n        if label is not None:\n            embed = self.embedding(label) #->(*,16ch)\n            outputs += torch.sum(embed*h,dim=1,keepdim=True) #->(*,1)\n        \n        outputs = torch.sigmoid(outputs)\n        return outputs","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:07.088462Z","iopub.execute_input":"2022-01-16T14:40:07.089043Z","iopub.status.idle":"2022-01-16T14:40:07.124994Z","shell.execute_reply.started":"2022-01-16T14:40:07.089000Z","shell.execute_reply":"2022-01-16T14:40:07.123934Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# #BigGAN\n# print(count_parameters(model=Generator(n_feat=27, codes_dim=24, n_classes=n_classes))) #z.shape=(*,120)\n# print(count_parameters(model=Discriminator(n_feat=33, n_classes=n_classes)))","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:07.126789Z","iopub.execute_input":"2022-01-16T14:40:07.127075Z","iopub.status.idle":"2022-01-16T14:40:07.143345Z","shell.execute_reply.started":"2022-01-16T14:40:07.127037Z","shell.execute_reply":"2022-01-16T14:40:07.142272Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"def generate_img(netG,fixed_noise,fixed_aux_labels=None):\n    if fixed_aux_labels is not None:\n        gen_image = netG(fixed_noise,fixed_aux_labels).to('cpu').clone().detach().squeeze(0)\n    else:\n        gen_image = netG(fixed_noise).to('cpu').clone().detach().squeeze(0)\n    #denormalize\n    gen_image = gen_image*0.5 + 0.5\n    gen_image_numpy = gen_image.numpy().transpose(0,2,3,1)\n    return gen_image_numpy\n\ndef show_generate_imgs(netG,fixed_noise,fixed_aux_labels=None):\n    gen_images_numpy = generate_img(netG,fixed_noise,fixed_aux_labels)\n\n    fig = plt.figure(figsize=(25, 16))\n    # display 10 images from each class\n    for i, img in enumerate(gen_images_numpy):\n        ax = fig.add_subplot(4, 8, i + 1, xticks=[], yticks=[])\n        plt.imshow(img)\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:07.145289Z","iopub.execute_input":"2022-01-16T14:40:07.145628Z","iopub.status.idle":"2022-01-16T14:40:07.157950Z","shell.execute_reply.started":"2022-01-16T14:40:07.145581Z","shell.execute_reply":"2022-01-16T14:40:07.156849Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#https://www.kaggle.com/osciiart/resnet34-mel-ver3-log-multi-hardaug?scriptVersionId=13887036\ndef cycle(iterable):\n    \"\"\"\n    dataloaderをiteratorに変換\n    :param iterable:\n    :return:\n    \"\"\"\n    while True:\n        for x in iterable:\n            yield x","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:07.159757Z","iopub.execute_input":"2022-01-16T14:40:07.160230Z","iopub.status.idle":"2022-01-16T14:40:07.170262Z","shell.execute_reply.started":"2022-01-16T14:40:07.160037Z","shell.execute_reply":"2022-01-16T14:40:07.169206Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#BigGAN\ndef run(lr_G=3e-4,lr_D=3e-4, beta1=0.0, beta2=0.999, nz=120, epochs=2, \n        n_ite_D=1, ema_decay_rate=0.999, show_epoch_list=None, output_freq=10):\n#     #G:4M, D:3M params\n#     netG = Generator(n_feat=22, codes_dim=24, n_classes=n_classes).to(device) #z.shape=(*,120)\n#     netD = Discriminator(n_feat=26, n_classes=n_classes).to(device)\n#     #G:6M, D:5M params\n#     netG = Generator(n_feat=27, codes_dim=24, n_classes=n_classes).to(device) #z.shape=(*,120)\n#     netD = Discriminator(n_feat=33, n_classes=n_classes).to(device)\n#     #G:8M, D:7M params\n#     netG = Generator(n_feat=32, codes_dim=24, n_classes=n_classes).to(device) #z.shape=(*,120)\n#     netD = Discriminator(n_feat=38, n_classes=n_classes).to(device)\n    #G:10M, D:8M params\n    netG = Generator(n_feat=36, codes_dim=24, n_classes=n_classes).to(device) #z.shape=(*,120)\n    netD = Discriminator(n_feat=42, n_classes=n_classes).to(device)\n#     #G:14M, D:12M params\n#     netG = Generator(n_feat=42, codes_dim=24, n_classes=n_classes).to(device) #z.shape=(*,120)\n#     netD = Discriminator(n_feat=50, n_classes=n_classes).to(device)\n#     #G:25M, D:22M params\n#     netG = Generator(n_feat=56, codes_dim=24, n_classes=n_classes).to(device) #z.shape=(*,120)\n#     netD = Discriminator(n_feat=68, n_classes=n_classes).to(device)\n#     #G:32M, D:31M params\n#     netG = Generator(n_feat=64, codes_dim=24, n_classes=n_classes).to(device) #z.shape=(*,120)\n#     netD = Discriminator(n_feat=80, n_classes=n_classes).to(device)\n\n    if EMA:\n        #EMA of G for sampling\n        netG_EMA = Generator(n_feat=42, codes_dim=24, n_classes=n_classes).to(device)\n        netG_EMA.load_state_dict(netG.state_dict())\n        for p in netG_EMA.parameters():\n            p.requires_grad = False\n\n        \n    print(count_parameters(netG))\n    print(count_parameters(netD))\n    \n    real_label = 0.9\n    fake_label = 0\n    \n    D_loss_list = []\n    G_loss_list = []\n    \n    dis_criterion = nn.BCELoss().to(device)\n\n    optimizerD = optim.Adam(netD.parameters(), lr=lr_D, betas=(beta1, beta2))\n    optimizerG = optim.Adam(netG.parameters(), lr=lr_G, betas=(beta1, beta2))\n    \n    fixed_noise = torch.randn(32, nz, 1, 1, device=device)\n    #fixed_noise = fixed_noise / fixed_noise.norm(dim=1, keepdim=True)\n    fixed_aux_labels     = np.random.randint(0,n_classes, 32)\n    fixed_aux_labels_ohe = np.eye(n_classes)[fixed_aux_labels]\n    fixed_aux_labels_ohe = torch.from_numpy(fixed_aux_labels_ohe[:,:,np.newaxis,np.newaxis])\n    fixed_aux_labels_ohe = fixed_aux_labels_ohe.float().to(device, non_blocking=True)\n\n    netG.train()\n    netD.train()\n\n    ### training here\n    for epoch in range(1,epochs+1):\n        if elapsed_time(start_time) > TIME_LIMIT:\n            print(f'elapsed_time go beyond {TIME_LIMIT} sec')\n            break\n        D_running_loss = 0\n        G_running_loss = 0\n        for ii, data in enumerate(train_loader):\n            ############################\n            # (1) Update D network\n            ###########################\n            for _ in range(n_ite_D):\n                \n                if LABEL_NOISE:\n                    real_label = 0.9\n                    fake_label = 0\n                    if np.random.random() < LABEL_NOISE_PROB:\n                        real_label = 0\n                        fake_label = 0.9\n                    \n                # train with real\n                netD.zero_grad()\n                real_images = data['img'].to(device, non_blocking=True) \n                batch_size  = real_images.size(0)\n                dis_labels  = torch.full((batch_size, 1), real_label, device=device) #shape=(*,)\n                aux_labels  = data['label'].long().to(device, non_blocking=True) #shape=(*,)\n                dis_output = netD(real_images, aux_labels) #dis shape=(*,1)\n                errD_real  = dis_criterion(dis_output, dis_labels)\n                errD_real.backward(retain_graph=True)\n\n                # train with fake\n                noise  = torch.randn(batch_size, nz, 1, 1, device=device)\n                #noise = noise / noise.norm(dim=1, keepdim=True)\n                aux_labels     = np.random.randint(0,n_classes, batch_size)\n                aux_labels_ohe = np.eye(n_classes)[aux_labels]\n                aux_labels_ohe = torch.from_numpy(aux_labels_ohe[:,:,np.newaxis,np.newaxis])\n                aux_labels_ohe = aux_labels_ohe.float().to(device, non_blocking=True)\n                aux_labels = torch.from_numpy(aux_labels).long().to(device, non_blocking=True)\n                \n                fake = netG(noise, aux_labels_ohe) #output shape=(*,3,64,64)\n                dis_labels.fill_(fake_label)\n                dis_output = netD(fake.detach(),aux_labels)\n                errD_fake  = dis_criterion(dis_output, dis_labels)\n                errD_fake.backward(retain_graph=True)\n                D_running_loss += (errD_real.item() + errD_fake.item())/len(train_loader)\n                optimizerD.step()\n\n            ############################\n            # (2) Update G network\n            ###########################\n            netG.zero_grad()\n            dis_labels.fill_(real_label)  # fake labels are real for generator cost\n            noise = torch.randn(batch_size, nz, 1, 1, device=device)\n            aux_labels     = np.random.randint(0,n_classes, batch_size)\n            aux_labels_ohe = np.eye(n_classes)[aux_labels]\n            aux_labels_ohe = torch.from_numpy(aux_labels_ohe[:,:,np.newaxis,np.newaxis])\n            aux_labels_ohe = aux_labels_ohe.float().to(device, non_blocking=True)\n            aux_labels = torch.from_numpy(aux_labels).long().to(device, non_blocking=True)\n            fake  = netG(noise, aux_labels_ohe)\n            \n            dis_output = netD(fake, aux_labels)\n            errG   = dis_criterion(dis_output, dis_labels)\n            errG.backward(retain_graph=True)\n            G_running_loss += errG.item()/len(train_loader)\n            optimizerG.step()\n        \n        if EMA:\n            #update netG_EMA\n            param_itr = cycle(netG.parameters())\n            for i,p_EMA in enumerate(netG_EMA.parameters()):\n                p = next(param_itr)\n                p_EMA.data = (1-ema_decay_rate)*p_EMA.data + ema_decay_rate*p.data\n                p_EMA.requires_grad = False\n        \n        #log\n        D_loss_list.append(D_running_loss)\n        G_loss_list.append(G_running_loss)\n        \n        #output\n        if epoch % output_freq == 0:\n            print('[{:d}/{:d}] D_loss = {:.3f}, G_loss = {:.3f}, elapsed_time = {:.1f} min'.format(epoch,epochs,D_running_loss,G_running_loss,elapsed_time(start_time)/60))\n            \n        if epoch in show_epoch_list:\n            print('epoch = {}'.format(epoch))\n            if not EMA:\n                show_generate_imgs(netG,fixed_noise,fixed_aux_labels_ohe)\n            elif EMA:\n                show_generate_imgs(netG_EMA,fixed_noise,fixed_aux_labels_ohe)\n            \n        if epoch % 100 == 0:\n            if not EMA:\n                torch.save(netG.state_dict(), f'generator_epoch{epoch}.pth')\n            elif EMA:\n                torch.save(netG_EMA.state_dict(), f'generator_epoch{epoch}.pth')\n    \n    if not EMA:\n        torch.save(netG.state_dict(), 'generator.pth')\n    elif EMA:\n        torch.save(netG_EMA.state_dict(), 'generator.pth')\n    torch.save(netD.state_dict(), 'discriminator.pth')\n    \n    res = {'netG':netG,\n           'netD':netD,\n           'nz':nz,\n           'fixed_noise':fixed_noise,\n           'fixed_aux_labels_ohe':fixed_aux_labels_ohe,\n           'D_loss_list':D_loss_list,\n           'G_loss_list':G_loss_list,\n          }\n    if EMA:\n        res['netG_EMA'] = netG_EMA\n        \n    return res","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:07.172369Z","iopub.execute_input":"2022-01-16T14:40:07.172757Z","iopub.status.idle":"2022-01-16T14:40:07.211209Z","shell.execute_reply.started":"2022-01-16T14:40:07.172715Z","shell.execute_reply":"2022-01-16T14:40:07.210014Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n#show_epoch_list = np.arange(0,100,1)\nshow_epoch_list = np.arange(0,500+10,10)\n\nres = run(lr_G=3e-4,lr_D=3e-4, beta1=0.0, beta2=0.999, nz=120, epochs=500, \n          n_ite_D=1, ema_decay_rate=None, show_epoch_list=show_epoch_list, output_freq=10)\n# res = run(lr_G=3e-4,lr_D=3e-4, beta1=0.5, beta2=0.999, nz=120, epochs=500, \n#           n_ite_D=1, ema_decay_rate=None, show_epoch_list=show_epoch_list, output_freq=10)","metadata":{"execution":{"iopub.status.busy":"2022-01-16T14:40:07.214371Z","iopub.execute_input":"2022-01-16T14:40:07.215164Z","iopub.status.idle":"2022-01-16T21:02:34.772520Z","shell.execute_reply.started":"2022-01-16T14:40:07.215119Z","shell.execute_reply":"2022-01-16T21:02:34.771267Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.plot(res['D_loss_list'], label='D_loss')\nplt.plot(res['G_loss_list'], label='G_loss')\nplt.grid()\nplt.legend()\nplt.title('loss history')","metadata":{"execution":{"iopub.status.busy":"2022-01-16T21:02:34.777361Z","iopub.execute_input":"2022-01-16T21:02:34.778290Z","iopub.status.idle":"2022-01-16T21:02:35.041159Z","shell.execute_reply.started":"2022-01-16T21:02:34.778242Z","shell.execute_reply":"2022-01-16T21:02:35.040105Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"code","source":"#truncation_trick\ndef submission_generate_images(res, truncated=None):\n    im_batch_size=50\n    n_images=10000\n    if not EMA:\n        netG = res['netG']\n    elif EMA:\n        netG = res['netG_EMA']\n    nz   = res['nz']\n    if not os.path.exists('../output_images'):\n        os.mkdir('../output_images')\n    for i_batch in range(0, n_images, im_batch_size):\n        if truncated is not None:\n            flag = True\n            while flag:\n                z = np.random.randn(100*im_batch_size*nz)\n                z = z[np.where(abs(z)<truncated)]\n                if len(z)>=im_batch_size*nz:\n                    flag=False\n            gen_z = torch.from_numpy(z[:im_batch_size*nz]).view(im_batch_size,nz,1,1)\n            gen_z = gen_z.float().to(device)\n        else:\n            gen_z = torch.randn(im_batch_size, nz, 1, 1, device=device)\n#         gen_z = gen_z / gen_z.norm(dim=1, keepdim=True)\n        aux_labels     = np.random.randint(0,n_classes, im_batch_size)\n        aux_labels_ohe = np.eye(n_classes)[aux_labels]\n        aux_labels_ohe = torch.from_numpy(aux_labels_ohe[:,:,np.newaxis,np.newaxis])\n        aux_labels_ohe = aux_labels_ohe.float().to(device)\n        \n        gen_images = netG(gen_z,aux_labels_ohe)\n        gen_images = gen_images.to(\"cpu\").clone().detach() #shape=(*,3,h,w), torch.Tensor\n        #denormalize\n        gen_images = gen_images*0.5 + 0.5\n        for i_image in range(gen_images.size(0)):\n            save_image(gen_images[i_image, :, :, :],\n                       os.path.join(f'../output_images', f'image_{i_batch+i_image:05d}.png'))\n    shutil.make_archive(f'images', 'zip', f'../output_images')","metadata":{"execution":{"iopub.status.busy":"2022-01-16T21:02:35.043007Z","iopub.execute_input":"2022-01-16T21:02:35.043342Z","iopub.status.idle":"2022-01-16T21:02:35.057778Z","shell.execute_reply.started":"2022-01-16T21:02:35.043296Z","shell.execute_reply":"2022-01-16T21:02:35.056326Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nsubmission_generate_images(res,truncated=0.8)","metadata":{"execution":{"iopub.status.busy":"2022-01-16T21:02:35.059503Z","iopub.execute_input":"2022-01-16T21:02:35.060499Z","iopub.status.idle":"2022-01-16T21:03:15.662952Z","shell.execute_reply.started":"2022-01-16T21:02:35.060445Z","shell.execute_reply":"2022-01-16T21:03:15.661873Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not EMA:\n    gen_image_numpy = generate_img(res['netG'],res['fixed_noise'],res['fixed_aux_labels_ohe'])\nelif EMA:\n    gen_image_numpy = generate_img(res['netG_EMA'],res['fixed_noise'],res['fixed_aux_labels_ohe'])\nfor img in gen_image_numpy:\n    plt.imshow(img)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2022-01-16T21:03:15.664699Z","iopub.execute_input":"2022-01-16T21:03:15.665334Z","iopub.status.idle":"2022-01-16T21:03:22.704141Z","shell.execute_reply.started":"2022-01-16T21:03:15.665268Z","shell.execute_reply":"2022-01-16T21:03:22.703060Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}