{"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":"code","source":"tpu_use = False\nif tpu_use == True:\n    !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n#     !python3 pytorch-xla-env-setup.py --version 1.9 --apt-packages libomp5 libopenblas-dev   \n#     !python3 pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev ","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:49.104213Z","iopub.execute_input":"2021-09-14T03:35:49.104663Z","iopub.status.idle":"2021-09-14T03:35:49.11775Z","shell.execute_reply.started":"2021-09-14T03:35:49.104548Z","shell.execute_reply":"2021-09-14T03:35:49.116839Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Give up for using TPU \n```\n!python3 pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev\n```\n```\nNumpy not Available\n```\n-> !python3 pytorch-xla-env-setup.py --version 1.7 or more\nis recommended\n\n\n```\n!python3 pytorch-xla-env-setup.py --version 1.9 --apt-packages libomp5 libopenblas-dev\n```\n```\nRuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:421 : Check failed: session->session()->Run(session_work->feed_inputs, session_work->outputs_handles, &outputs) == ::tensorflow::Status::OK() (Aborted: Session bae0b7f7935a353c is not found. vs. OK)\n*** Begin stack trace ***\n\ttensorflow::CurrentStackTrace()\n\t\n\txla::util::MultiWait::Complete(std::function<void ()> const&)\n```\n-> !python3 pytorch-xla-env-setup.py --version nightly\nis recommended\n\nLoop","metadata":{}},{"cell_type":"code","source":"import sys\nimport glob\nimport os\nimport cv2\nfrom pathlib import Path\nimport random\nimport time\nfrom tqdm.notebook import tqdm\nimport itertools\nimport shutil\n\nimport numpy as np \nimport pandas as pd \n\nimport PIL.Image\nfrom IPython.display import Image\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.nn.init as init\nfrom torch.utils.data import random_split, DataLoader\nimport torchvision.models as models\nimport torchvision.transforms as transforms","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:49.119492Z","iopub.execute_input":"2021-09-14T03:35:49.120134Z","iopub.status.idle":"2021-09-14T03:35:50.701545Z","shell.execute_reply.started":"2021-09-14T03:35:49.120092Z","shell.execute_reply":"2021-09-14T03:35:50.700679Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if tpu_use == True:\n    import torch_xla\n    import torch_xla.core.xla_model as xm\n    import torch_xla.distributed.xla_multiprocessing as xmp\n    import torch_xla.distributed.parallel_loader as pl\n    import torch_xla.utils.serialization as xser\n    device = xm.xla_device()\nelse:\n    device = 'cuda' if torch.cuda.is_available() else 'cpu'\ndevice","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:50.703211Z","iopub.execute_input":"2021-09-14T03:35:50.70353Z","iopub.status.idle":"2021-09-14T03:35:50.772722Z","shell.execute_reply.started":"2021-09-14T03:35:50.703495Z","shell.execute_reply":"2021-09-14T03:35:50.771874Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load Data","metadata":{}},{"cell_type":"code","source":"def get_data_as_nparray(dir):\n    data = []\n    for i, image_path in enumerate(glob.glob(os.path.join(dir,'*.jpg'))):\n        if i > 5:\n            continue\n        img = cv2.imread(image_path) # (height,width,channels)\n        img_expanded = np.expand_dims(img,axis=0) # (1,height,width,channels)\n        data.append(img_expanded)\n    # (n_samples,height,width,channels)\n    data_np = np.concatenate(data,axis=0)\n    return data_np\n\ndef show_imgs(imd_data,height,width):\n    %matplotlib inline\n    # Show multi img\n    plt.figure(figsize=(12,12))\n    for i,d in enumerate(imd_data):\n        if i >= width*height:         \n            continue\n        plt.subplot(height,width,i+1)\n        plt.imshow(ｄ)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:50.774259Z","iopub.execute_input":"2021-09-14T03:35:50.774948Z","iopub.status.idle":"2021-09-14T03:35:50.789117Z","shell.execute_reply.started":"2021-09-14T03:35:50.774908Z","shell.execute_reply":"2021-09-14T03:35:50.788248Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"style_dir = '/kaggle/input/gan-getting-started/monet_jpg/'\nstyle_data = get_data_as_nparray(style_dir)\nshow_imgs(style_data,1,3)\nprint(style_data.shape)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:50.790502Z","iopub.execute_input":"2021-09-14T03:35:50.790991Z","iopub.status.idle":"2021-09-14T03:35:51.416881Z","shell.execute_reply.started":"2021-09-14T03:35:50.79083Z","shell.execute_reply":"2021-09-14T03:35:51.416046Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"photo_dir = '/kaggle/input/gan-getting-started/photo_jpg/'\nphoto_data = get_data_as_nparray(photo_dir)\nshow_imgs(photo_data,1,3)\n","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:51.418036Z","iopub.execute_input":"2021-09-14T03:35:51.418382Z","iopub.status.idle":"2021-09-14T03:35:52.392427Z","shell.execute_reply.started":"2021-09-14T03:35:51.41834Z","shell.execute_reply":"2021-09-14T03:35:52.391434Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## random,torch　ensuring reproducibility\n[reference](https://www.google.com/search?q=++++torch.backends.cudnn.deterministic+qiita&sxsrf=AOaemvKxBaC4koa_BLVExQIaGCbCwYaxHg%3A1630678601472&ei=SS4yYdejHI7_0ASTyIuQAg&oq=++++torch.backends.cudnn.deterministic+qiita&gs_lcp=Cgdnd3Mtd2l6EAM6BwgjELADECdKBAhBGAFQnBZYnBZgyx5oAXAAeACAAT2IAT2SAQExmAEAoAECoAEByAEBwAEB&sclient=gws-wiz&ved=0ahUKEwiX7KGW_-LyAhWOP5QKHRPkAiIQ4dUDCA4&uact=5)","metadata":{}},{"cell_type":"code","source":"def set_seed(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)   # pythonのハッシュベースの操作の再現性を担保\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.backends.cudnn.benchmark = False      # False:再現性確保 True：速度確保\n    torch.backends.cudnn.deterministic = True   # PyTorchの操作の中には非決定的なものがあります。それらを決定的なものにします。","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.39481Z","iopub.execute_input":"2021-09-14T03:35:52.39517Z","iopub.status.idle":"2021-09-14T03:35:52.400521Z","shell.execute_reply.started":"2021-09-14T03:35:52.395132Z","shell.execute_reply":"2021-09-14T03:35:52.399336Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed = 95\nset_seed(seed)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.40273Z","iopub.execute_input":"2021-09-14T03:35:52.403136Z","iopub.status.idle":"2021-09-14T03:35:52.413886Z","shell.execute_reply.started":"2021-09-14T03:35:52.403099Z","shell.execute_reply":"2021-09-14T03:35:52.413065Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Create Dataset with Shuffled Pair\n[reference](https://qiita.com/kumonk/items/0f3cad018cc9aec67a63)","metadata":{}},{"cell_type":"code","source":"class MyDataset(torch.utils.data.Dataset):\n\n    def __init__(self, style_dir, photo_dir, imageSize=(256,256), transform=None, normalize=True, diffPairAugmentRate=5):\n        if normalize:\n            self.transform = transforms.Compose([\n                transforms.Resize(imageSize), \n                transforms.ToTensor(), \n                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n            ])\n        else:\n            self.transform = transforms.Compose([\n                transforms.Resize(imageSize), \n                transforms.ToTensor(), \n            ])\n        # Dataのパスリスト\n        self.style_paths = [str(p) for p in Path(style_dir).glob(\"*.jpg\")]*diffPairAugmentRate\n        self.photo_paths = [str(p) for p in Path(photo_dir).glob(\"*.jpg\")]*diffPairAugmentRate\n        # Shuffle Pairs of input and grandtruth data\n        random.shuffle(self.style_paths)\n        random.shuffle(self.photo_paths)\n        self.data_num = min(len(self.style_paths),len(self.photo_paths)) # ここが__len__の返り値になる\n\n\n    def __getitem__(self, idx):\n        s = self.style_paths[idx]\n        style_img = PIL.Image.open(s)\n        p = self.photo_paths[idx]\n        photo_img = PIL.Image.open(p)\n\n        if self.transform:\n            out_style = self.transform(style_img)\n            out_photo = self.transform(photo_img)            \n        \n        # ForLabel\n        # out_label = p.split(\"\\\\\")\n        # out_label = self.class_to_idx[out_label[3]]\n\n        return out_photo, out_style,\n        \n    def __len__(self):\n        return self.data_num","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.415328Z","iopub.execute_input":"2021-09-14T03:35:52.415736Z","iopub.status.idle":"2021-09-14T03:35:52.428383Z","shell.execute_reply.started":"2021-09-14T03:35:52.415701Z","shell.execute_reply":"2021-09-14T03:35:52.427476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# dataset = MyDataset(style_dir, photo_dir, normalize=False)\ndiffPairAugmentRate=10\ndataset = MyDataset(style_dir, photo_dir, normalize=True,diffPairAugmentRate=diffPairAugmentRate)\ntrain_size = int(len(dataset) * 0.95) \nval_size = len(dataset) - train_size \ntrain_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.429828Z","iopub.execute_input":"2021-09-14T03:35:52.430326Z","iopub.status.idle":"2021-09-14T03:35:52.523245Z","shell.execute_reply.started":"2021-09-14T03:35:52.43029Z","shell.execute_reply":"2021-09-14T03:35:52.52246Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !python -m pip install numpy==1.18.3\n# import numpy as np\n# style_img, photo_img = next(iter(dataset))\n\n# print(len(style_img))\n# print(style_img.size())\n# show_imgs([style_img.permute(1, 2, 0),photo_img.permute(1, 2, 0)],1,2)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.524397Z","iopub.execute_input":"2021-09-14T03:35:52.524753Z","iopub.status.idle":"2021-09-14T03:35:52.528627Z","shell.execute_reply.started":"2021-09-14T03:35:52.524719Z","shell.execute_reply":"2021-09-14T03:35:52.527324Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def unnorm(img, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):\n    for t, m, s in zip(img, mean, std):\n        t.mul_(s).add_(s)  \n    return img\n\n# show_imgs([unnorm(style_img).permute(1, 2, 0),unnorm(photo_img).permute(1, 2, 0)],1,2)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.53008Z","iopub.execute_input":"2021-09-14T03:35:52.530539Z","iopub.status.idle":"2021-09-14T03:35:52.539838Z","shell.execute_reply.started":"2021-09-14T03:35:52.530506Z","shell.execute_reply":"2021-09-14T03:35:52.53899Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## DataLoader ensuring reproducibility\n[reference](https://www.google.com/search?q=++++torch.backends.cudnn.deterministic+qiita&sxsrf=AOaemvKxBaC4koa_BLVExQIaGCbCwYaxHg%3A1630678601472&ei=SS4yYdejHI7_0ASTyIuQAg&oq=++++torch.backends.cudnn.deterministic+qiita&gs_lcp=Cgdnd3Mtd2l6EAM6BwgjELADECdKBAhBGAFQnBZYnBZgyx5oAXAAeACAAT2IAT2SAQExmAEAoAECoAEByAEBwAEB&sclient=gws-wiz&ved=0ahUKEwiX7KGW_-LyAhWOP5QKHRPkAiIQ4dUDCA4&uact=5)\n## DataLoader speedy\n```\nnum_workers=os.cpu_count()\npin_memory=True\n```\n[reference](https://qiita.com/sugulu_Ogawa_ISID/items/62f5f7adee083d96a587#12-pin_memory)","metadata":{}},{"cell_type":"code","source":"def seed_worker(worker_id):\n    worker_seed = torch.initial_seed() % 2**32\n    np.random.seed(worker_seed)\n    random.seed(worker_seed)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.541177Z","iopub.execute_input":"2021-09-14T03:35:52.541599Z","iopub.status.idle":"2021-09-14T03:35:52.548617Z","shell.execute_reply.started":"2021-09-14T03:35:52.541549Z","shell.execute_reply":"2021-09-14T03:35:52.547643Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Checkpoint Load and Save","metadata":{}},{"cell_type":"code","source":"def load_checkpoint(ckpt_path, map_location=None):\n    ckpt = torch.load(ckpt_path, map_location=map_location)\n    print(' [*] Loading checkpoint from %s succeed!' % ckpt_path)\n    return ckpt\ndef save_checkpoint(state, save_path):\n    if tpu_use==True:\n        xm.save(state, save_path)\n    else:\n        torch.save(state, save_path)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.549896Z","iopub.execute_input":"2021-09-14T03:35:52.550421Z","iopub.status.idle":"2021-09-14T03:35:52.558321Z","shell.execute_reply.started":"2021-09-14T03:35:52.550381Z","shell.execute_reply":"2021-09-14T03:35:52.557439Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model\n\n[ReflectionPad2d](https://teratail.com/questions/254795)","metadata":{}},{"cell_type":"code","source":"def Upsample(in_ch, out_ch, use_dropout=True, dropout_ratio=0.1):\n    if use_dropout:\n        return nn.Sequential(\n            nn.ConvTranspose2d(in_ch, out_ch, 3, stride=2, padding=1, output_padding=1),\n            nn.InstanceNorm2d(out_ch),\n            nn.Dropout(dropout_ratio),\n            nn.GELU()\n        )\n    else:\n        return nn.Sequential(\n            nn.ConvTranspose2d(in_ch, out_ch, 3, stride=2, padding=1, output_padding=1),\n            nn.InstanceNorm2d(out_ch),\n            nn.GELU()\n        )\n    \ndef Convlayer(in_ch, out_ch, kernel_size=3, stride=1, use_leaky=True, use_inst_norm=True, use_pad=True):\n    if use_pad:\n        conv = nn.Conv2d(in_ch, out_ch, kernel_size, stride, 1, bias=True)\n    else:\n        conv = nn.Conv2d(in_ch, out_ch, kernel_size, stride, 0, bias=True)\n\n    if use_leaky:\n        actv = nn.LeakyReLU(negative_slope=0.2, inplace=True)\n    else:\n        actv = nn.GELU()\n\n    if use_inst_norm:\n        norm = nn.InstanceNorm2d(out_ch)\n    else:\n        norm = nn.BatchNorm2d(out_ch)\n\n    return nn.Sequential(\n        conv,\n        norm,\n        actv\n    )\n\nclass Resblock(nn.Module):\n    def __init__(self, in_features, use_dropout=True, dropout_ratio=0.1):\n        super().__init__()\n        layers = list()\n        layers.append(nn.ReflectionPad2d(1))\n        layers.append(Convlayer(in_features, in_features, 3, 1, use_leaky=True, use_inst_norm=True, use_pad=False))\n        layers.append(nn.Dropout(dropout_ratio))\n        layers.append(nn.ReflectionPad2d(1))\n        layers.append(nn.Conv2d(in_features, in_features, 3, 1, padding=0, bias=True))\n        layers.append(nn.InstanceNorm2d(in_features))\n        self.res = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return x + self.res(x)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.559994Z","iopub.execute_input":"2021-09-14T03:35:52.560301Z","iopub.status.idle":"2021-09-14T03:35:52.576335Z","shell.execute_reply.started":"2021-09-14T03:35:52.560276Z","shell.execute_reply":"2021-09-14T03:35:52.575459Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, in_ch, out_ch, num_res_blocks=6):\n        super().__init__()\n        model = list()\n        model.append(nn.ReflectionPad2d(3))\n        model.append(Convlayer(in_ch, 64, kernel_size=7, stride=1, use_leaky=True, use_inst_norm=True, use_pad=False))\n        model.append(Convlayer(64, 128, kernel_size=3, stride=2, use_leaky=True, use_inst_norm=True, use_pad=True))\n        model.append(Convlayer(128, 256, kernel_size=3, stride=2, use_leaky=True, use_inst_norm=True, use_pad=True))\n        for _ in range(num_res_blocks):\n            model.append(Resblock(256))\n        model.append(Upsample(256, 128))\n        model.append(Upsample(128, 64))\n        model.append(nn.ReflectionPad2d(3))\n        model.append(nn.Conv2d(64, out_ch, kernel_size=7, stride=1, padding=0))\n        model.append(nn.Tanh())\n\n        self.gen = nn.Sequential(*model)\n\n    def forward(self, x):\n        return self.gen(x)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.577595Z","iopub.execute_input":"2021-09-14T03:35:52.577954Z","iopub.status.idle":"2021-09-14T03:35:52.589297Z","shell.execute_reply.started":"2021-09-14T03:35:52.577915Z","shell.execute_reply":"2021-09-14T03:35:52.588471Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, in_ch, num_layers=4):\n        super().__init__()\n        model = list()\n        model.append(nn.Conv2d(in_ch, 64, kernel_size=4, stride=2, padding=1))\n        model.append(nn.LeakyReLU(negative_slope=0.2, inplace=True))\n        for i in range(1, num_layers):\n            in_chs = 64 * 2**(i-1)\n            out_chs = in_chs * 2\n            if i == num_layers -1:\n                model.append(Convlayer(in_chs, out_chs, kernel_size=4, stride=1))\n            else:\n                model.append(Convlayer(in_chs, out_chs, kernel_size=4, stride=2))\n        model.append(nn.Conv2d(out_chs, 1, kernel_size=4, stride=1, padding=1))\n#         model.append(nn.Sigmoid())\n        self.disc = nn.Sequential(*model)\n\n    def forward(self, x):\n        return self.disc(x)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.590508Z","iopub.execute_input":"2021-09-14T03:35:52.590858Z","iopub.status.idle":"2021-09-14T03:35:52.601295Z","shell.execute_reply.started":"2021-09-14T03:35:52.590821Z","shell.execute_reply":"2021-09-14T03:35:52.60046Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# def init_weights(net, init_type='normal', gain=0.02):\n#     def init_func(m):\n#         classname = m.__class__.__name__\n#         if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):\n#             init.normal_(m.weight.data, 0.0, gain)\n#             if hasattr(m, 'bias') and m.bias is not None:\n#                 init.constant_(m.bias.data, 0.0)\n#         elif classname.find('BatchNorm2d') != -1:\n#             init.normal_(m.weight.data, 1.0, gain)\n#             init.constant_(m.bias.data, 0.0)\n#     net.apply(init_func)\ndef init_weights(m, init_type='normal', gain=0.02):\n    def init_func(m):\n        classname = m.__class__.__name__\n        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):\n            init.normal_(m.weight.data, 0.0, gain)\n            if hasattr(m, 'bias') and m.bias is not None:\n                init.constant_(m.bias.data, 0.0)\n        elif classname.find('BatchNorm2d') != -1:\n            init.normal_(m.weight.data, 1.0, gain)\n            init.constant_(m.bias.data, 0.0)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.60276Z","iopub.execute_input":"2021-09-14T03:35:52.603157Z","iopub.status.idle":"2021-09-14T03:35:52.610749Z","shell.execute_reply.started":"2021-09-14T03:35:52.603124Z","shell.execute_reply":"2021-09-14T03:35:52.609853Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def update_req_grad(models, requires_grad=True):\n    for model in models:\n        for param in model.parameters():\n            param.requires_grad = requires_grad\n            \n# https://arxiv.org/pdf/1612.07828.pdf\n# Save 50 generated fake imgs and sample through them\n# to feed discriminators to avoid large oscillations \n# from iterations to iterations.\nclass sample_fake(object):\n    def __init__(self, max_imgs=50):\n        self.max_imgs = max_imgs\n        self.cur_img = 0\n        self.imgs = list()\n\n    def __call__(self, imgs):\n        ret = list()\n        for img in imgs:\n            if self.cur_img < self.max_imgs:\n                self.imgs.append(img)\n                ret.append(img)\n                self.cur_img += 1\n            else:\n                if np.random.ranf() > 0.5:\n                    idx = np.random.randint(0, self.max_imgs)\n                    ret.append(self.imgs[idx])\n                    self.imgs[idx] = img\n                else:\n                    ret.append(img)\n        return ret\n    \nclass lr_sched():\n    def __init__(self, decay_epochs=20, total_epochs=40):\n        self.decay_epochs = decay_epochs\n        self.total_epochs = total_epochs\n\n    def step(self, epoch_num):\n        if epoch_num <= self.decay_epochs:\n            return 1.0\n        else:\n            fract = (epoch_num - self.decay_epochs)  / (self.total_epochs - self.decay_epochs)\n            return 1.0 - fract\n\nclass AvgStats(object):\n    def __init__(self):\n        self.reset()\n        \n    def reset(self):\n        self.losses =[]\n        self.its = []\n        \n    def append(self, loss, it):\n        self.losses.append(loss)\n        self.its.append(it)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.614227Z","iopub.execute_input":"2021-09-14T03:35:52.6148Z","iopub.status.idle":"2021-09-14T03:35:52.627024Z","shell.execute_reply.started":"2021-09-14T03:35:52.614764Z","shell.execute_reply":"2021-09-14T03:35:52.626174Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# GAN","metadata":{}},{"cell_type":"code","source":"class CycleGAN(object):\n    def __init__(self, in_ch, out_ch, epochs, device, start_lr=5e-4, lmbda=10, idt_coef=0.5, decay_epoch=0):\n        self.epochs = epochs\n        self.decay_epoch = decay_epoch if decay_epoch > 0 else int(self.epochs/2)\n        self.lmbda = lmbda\n        self.idt_coef = idt_coef\n        self.device = device\n        if tpu_use == True:\n            self.gen_s2p = xmp.MpModelWrapper(Generator(in_ch, out_ch))\n            self.gen_p2s = xmp.MpModelWrapper(Generator(in_ch, out_ch))\n            self.desc_s = xmp.MpModelWrapper(Discriminator(in_ch))\n            self.desc_p = xmp.MpModelWrapper(Discriminator(in_ch))\n        else:   \n            self.gen_s2p = Generator(in_ch, out_ch)\n            self.gen_p2s = Generator(in_ch, out_ch)\n            self.desc_s = Discriminator(in_ch)\n            self.desc_p = Discriminator(in_ch)\n        self.init_models()\n        self.mse_loss = nn.MSELoss()\n        self.l1_loss = nn.L1Loss()\n#         self.adam_gen = torch.optim.Adam(itertools.chain(self.gen_s2p.parameters(), self.gen_p2s.parameters()),\n#                                          lr= start_lr, betas=(0.5, 0.999))\n#         self.adam_desc = torch.optim.Adam(itertools.chain(self.desc_s.parameters(), self.desc_p.parameters()),\n#                                           lr=start_lr, betas=(0.5, 0.999))\n        self.adam_gen = torch.optim.AdamW(itertools.chain(self.gen_s2p.parameters(), self.gen_p2s.parameters()),\n                                         lr= start_lr)\n        self.adam_desc = torch.optim.AdamW(itertools.chain(self.desc_s.parameters(), self.desc_p.parameters()),\n                                          lr=start_lr)\n        self.sample_style = sample_fake()\n        self.sample_photo = sample_fake()\n        gen_lr = lr_sched(self.decay_epoch, self.epochs)\n        desc_lr = lr_sched(self.decay_epoch, self.epochs)\n        self.gen_lr_sched = torch.optim.lr_scheduler.LambdaLR(self.adam_gen, gen_lr.step)\n        self.desc_lr_sched = torch.optim.lr_scheduler.LambdaLR(self.adam_desc, desc_lr.step)\n        self.gen_stats = AvgStats()\n        self.desc_stats = AvgStats()\n        self.val_gen_stats =  AvgStats()\n        self.val_desc_stats =  AvgStats()\n        \n    def init_models(self):\n        init_weights(self.gen_s2p)\n        init_weights(self.gen_p2s)\n        init_weights(self.desc_s)\n        init_weights(self.desc_p)\n        self.gen_s2p = self.gen_s2p.to(self.device)\n        self.gen_p2s = self.gen_p2s.to(self.device)\n        self.desc_s = self.desc_s.to(self.device)\n        self.desc_p = self.desc_p.to(self.device)\n    \n    def make_dataloader(self,train_dataset, val_dataset):\n        if tpu_use==True:\n            train_sampler = DistributedSampler(\n                train_dataset,\n                num_replicas=xm.xrt_world_size(),\n                rank=xm.get_ordinal(),\n                shuffle=True\n                )\n            val_sampler = DistributedSampler(\n                val_dataset,\n                num_replicas=xm.xrt_world_size(),\n                rank=xm.get_ordinal(),\n                shuffle=True\n                )\n            train_sampler.set_epoch(self.epochs)\n            val_sampler.set_epoch(self.epochs)\n\n#             g = torch.Generator()\n#             g.manual_seed(seed)\n            img_train = DataLoader(train_dataset,\n                                    batch_size=1,\n                                    sampler=train_sampler, \n                                    num_workers=os.cpu_count(),\n                                    worker_init_fn=seed_worker,\n                                    generator=g,\n                                    pin_memory=True)\n            img_val = DataLoader(val_dataset,\n                                    batch_size=1,\n                                    sampler=val_sampler, \n                                    num_workers=os.cpu_count(),\n                                    worker_init_fn=seed_worker,\n                                    generator=g,\n                                    pin_memory=True)\n            img_train = pl.ParallelLoader(img_train, [device]).per_device_loader(device)\n            img_val = pl.ParallelLoader(img_val, [device]).per_device_loader(device)\n        \n        \n        else:  \n            img_train = DataLoader(train_dataset,\n                                    batch_size=1,\n                                    num_workers=os.cpu_count(),\n                                    worker_init_fn=seed_worker,\n#                                     generator=g,\n                                    pin_memory=True)\n            img_val = DataLoader(val_dataset,\n                                    batch_size=1,\n                                    num_workers=os.cpu_count(),\n                                    worker_init_fn=seed_worker,\n#                                     generator=g,\n                                    pin_memory=True) \n            return img_train, img_val\n    \n    def run_gen(self,photo_img, style_img):\n        # Forward pass through generator\n        fake_photo = self.gen_s2p(style_img)\n        fake_style = self.gen_p2s(photo_img)\n\n        cycl_style = self.gen_p2s(fake_photo)\n        cycl_photo = self.gen_s2p(fake_style)\n\n        id_style = self.gen_p2s(style_img)\n        id_photo = self.gen_s2p(photo_img)\n\n        # generator losses - identity, Adversarial, cycle consistency\n        idt_loss_style = self.l1_loss(id_style, style_img) * self.lmbda * self.idt_coef\n        idt_loss_photo = self.l1_loss(id_photo, photo_img) * self.lmbda * self.idt_coef\n\n        cycle_loss_style = self.l1_loss(cycl_style, style_img) * self.lmbda\n        cycle_loss_photo = self.l1_loss(cycl_photo, photo_img) * self.lmbda\n\n        style_desc = self.desc_s(fake_style)\n        photo_desc = self.desc_p(fake_photo)\n\n        real = torch.ones(style_desc.size()).to(self.device)\n\n        adv_loss_style = self.mse_loss(style_desc, real)\n        adv_loss_photo = self.mse_loss(photo_desc, real)\n\n        # total generator loss\n        total_gen_loss = cycle_loss_style + adv_loss_style\\\n                      + cycle_loss_photo + adv_loss_photo\\\n                      + idt_loss_style + idt_loss_photo\n        return total_gen_loss, fake_photo, fake_style\n    \n    def run_desc(self, photo_img, style_img, fake_style, fake_photo):\n        fake_style = torch.tensor(fake_style).to(self.device)\n        fake_photo = torch.tensor(fake_photo).to(self.device)\n\n        style_desc_real = self.desc_s(style_img)\n        style_desc_fake = self.desc_s(fake_style)\n        photo_desc_real = self.desc_p(photo_img)\n        photo_desc_fake = self.desc_p(fake_photo)\n\n        real = torch.ones(style_desc_real.size()).to(self.device)\n        fake = torch.zeros(style_desc_fake.size()).to(self.device)\n\n        # Descriminator losses\n        # --------------------\n        style_desc_real_loss = self.mse_loss(style_desc_real, real)\n        style_desc_fake_loss = self.mse_loss(style_desc_fake, fake)\n        photo_desc_real_loss = self.mse_loss(photo_desc_real, real)\n        photo_desc_fake_loss = self.mse_loss(photo_desc_fake, fake)\n\n        style_desc_loss = (style_desc_real_loss + style_desc_fake_loss) / 2\n        photo_desc_loss = (photo_desc_real_loss + photo_desc_fake_loss) / 2\n        total_desc_loss = style_desc_loss + photo_desc_loss\n        \n        return total_desc_loss, style_desc_loss, photo_desc_loss\n    \n    def train(self, train_dataset,val_dataset,diffPairAugmentRate):\n        best_val=99999\n        img_train, img_val = self.make_dataloader(train_dataset, val_dataset)\n        for epoch in range(self.epochs):\n            \n            start_time = time.time()\n            avg_gen_loss = 0.0\n            avg_desc_loss = 0.0\n            val_avg_gen_loss = 0.0\n            val_avg_desc_loss = 0.0\n            \n            # For cut augmented dataset\n            original_len_val = img_val.__len__() / diffPairAugmentRate\n            idx_max_val = original_len_val  * (diffPairAugmentRate-1) - 1 \n                \n            original_len_train = img_train.__len__() / diffPairAugmentRate\n            idx_max_train = original_len_train * (diffPairAugmentRate-1) - 1 \n            \n            idx_start_val = random.randint(0,idx_max_val)\n            idx_start_train = random.randint(0,idx_max_train)\n            \n            # Validation\n            self.gen_s2p.eval()\n            self.gen_p2s.eval()\n            self.desc_s.eval()\n            self.desc_p.eval()\n\n            for i, (photo_real, style_real) in enumerate(img_val):\n                if i < idx_start_val:\n                    continue\n                elif i >= idx_start_val + original_len_val:\n                    continue\n                    \n                with torch.no_grad():\n                    # Generator\n                    photo_img, style_img = photo_real.to(device), style_real.to(device)\n                    total_gen_loss, fake_photo, fake_style= self.run_gen(photo_img, style_img)\n                    # Descriminator\n                    fake_style = self.sample_style([fake_style.cpu().data.numpy()])[0]\n                    fake_photo = self.sample_photo([fake_photo.cpu().data.numpy()])[0]\n                    total_desc_loss, style_desc_loss, photo_desc_loss = self.run_desc(photo_img, style_img, fake_style, fake_photo)                 \n                    \n                    val_avg_gen_loss += total_gen_loss.item()\n                    val_avg_desc_loss += total_desc_loss.item()\n                    \n            # Train\n            t = tqdm(img_train, leave=False, total=img_train.__len__())\n            self.gen_s2p.train()\n            self.gen_p2s.train()\n            self.desc_s.train()\n            self.desc_p.train()\n            for i, (photo_real, style_real) in enumerate(t):\n                if i < idx_start_train:\n                    continue\n                elif i >= idx_start_train + original_len_train:\n                    continue\n                photo_img, style_img = photo_real.to(device), style_real.to(device)\n                \n                # Generator\n                update_req_grad([self.desc_s, self.desc_p], False)\n                self.adam_gen.zero_grad()\n                total_gen_loss, fake_photo, fake_style= self.run_gen(photo_img, style_img)\n                avg_gen_loss += total_gen_loss.item()\n                total_gen_loss.backward()\n                self.adam_gen.step()\n\n                # Descriminator\n                update_req_grad([self.desc_s, self.desc_p], True)\n                self.adam_desc.zero_grad()\n                fake_style = self.sample_style([fake_style.cpu().data.numpy()])[0]\n                fake_photo = self.sample_photo([fake_photo.cpu().data.numpy()])[0]\n                total_desc_loss, style_desc_loss, photo_desc_loss = self.run_desc(photo_img, style_img, fake_style, fake_photo)\n                avg_desc_loss += total_desc_loss.item()\n\n                # Backward\n                style_desc_loss.backward()\n                photo_desc_loss.backward()\n                self.adam_desc.step()\n                \n                t.set_postfix(gen_loss=total_gen_loss.item(), desc_loss=total_desc_loss.item())\n            \n            val_avg_gen_loss /= img_val.__len__()\n#             val_avg_desc_loss /= img_val.__len__()\n            val_avg_desc_loss /= img_val.__len__()*0.1 # for easy-look graph\n           \n            avg_gen_loss /= img_train.__len__()\n#             avg_desc_loss /= img_train.__len__()\n            avg_desc_loss /= img_train.__len__()*0.1 # for easy-look graph\n            time_req = time.time() - start_time\n\n            val_loss = val_avg_gen_loss+val_avg_desc_loss\n            if best_val >= val_loss:\n    \n                save_dict = {\n                    'epoch': epoch+1,\n                    'gen_s2p': gan.gen_s2p.state_dict(),\n                    'gen_p2s': gan.gen_p2s.state_dict(),\n                    'desc_s': gan.desc_s.state_dict(),\n                    'desc_p': gan.desc_p.state_dict(),\n                    'optimizer_gen': gan.adam_gen.state_dict(),\n                    'optimizer_desc': gan.adam_desc.state_dict()\n                }\n                save_checkpoint(save_dict, 'current.ckpt')\n\n            self.val_gen_stats.append(val_avg_gen_loss, time_req)\n            self.val_desc_stats.append(val_avg_desc_loss, time_req)\n            self.gen_stats.append(avg_gen_loss, time_req)\n            self.desc_stats.append(avg_desc_loss, time_req)\n            \n            if tpu_use == True:    \n                xm.master_print(\"Epoch: (%d) | Gen ValLoss:%f | Disc ValLoss:%f \" % \n                        (epoch+0, val_avg_gen_loss,val_avg_desc_loss))\n                xm.master_print(\"Epoch: (%d) | Gen TraLoss:%f | Disc TraLoss:%f \" % \n                        (epoch+1, avg_gen_loss, avg_desc_loss))\n            else:\n                print(\"Epoch: (%d) | Gen ValLoss:%f | Disc ValLoss:%f \" % \n                        (epoch+0, val_avg_gen_loss,val_avg_desc_loss))\n                print(\"Epoch: (%d) | Gen TraLoss:%f | Disc TraLoss:%f \" % \n                        (epoch+1, avg_gen_loss, avg_desc_loss))\n            self.gen_lr_sched.step()\n            self.desc_lr_sched.step()","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.629019Z","iopub.execute_input":"2021-09-14T03:35:52.629492Z","iopub.status.idle":"2021-09-14T03:35:52.675634Z","shell.execute_reply.started":"2021-09-14T03:35:52.629455Z","shell.execute_reply":"2021-09-14T03:35:52.674896Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gan = CycleGAN(3, 3, 100, device)\nif tpu_use == True:\n    xmp.spawn(gan.train(train_dataset,val_dataset), nprocs=xm.xrt_world_size(), start_method='fork')\nelse:\n    gan.train(train_dataset,val_dataset,diffPairAugmentRate)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:35:52.677105Z","iopub.execute_input":"2021-09-14T03:35:52.677454Z","iopub.status.idle":"2021-09-14T03:36:10.716136Z","shell.execute_reply.started":"2021-09-14T03:35:52.67742Z","shell.execute_reply":"2021-09-14T03:36:10.713412Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.xlabel(\"Epochs\")\nplt.ylabel(\"Losses\")\nplt.plot(gan.gen_stats.losses, 'r', label='Generator Loss')\nplt.plot(gan.desc_stats.losses, 'b', label='Descriminator Loss')\nplt.plot(gan.val_gen_stats.losses, 'magenta', label='Generator Loss')\nplt.plot(gan.val_desc_stats.losses, 'cyan', label='Descriminator Loss')\nplt.legend()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:36:14.9059Z","iopub.execute_input":"2021-09-14T03:36:14.90625Z","iopub.status.idle":"2021-09-14T03:36:15.060271Z","shell.execute_reply.started":"2021-09-14T03:36:14.906219Z","shell.execute_reply":"2021-09-14T03:36:15.059499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generation for samples by best model","metadata":{}},{"cell_type":"code","source":"ckpt_path = 'current.ckpt'\nckpt= load_checkpoint(ckpt_path)\nprint(ckpt.keys())\ngan = CycleGAN(3, 3, 50, device)\ngan.gen_s2p.load_state_dict(ckpt['gen_s2p'])\ngan.gen_p2s.load_state_dict(ckpt['gen_p2s'])\ngan.desc_s.load_state_dict(ckpt['desc_s'])\ngan.desc_p.load_state_dict(ckpt['desc_p'])\ngan.adam_gen.load_state_dict(ckpt['optimizer_gen'])\ngan.adam_desc.load_state_dict(ckpt['optimizer_desc'])","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:36:18.523404Z","iopub.execute_input":"2021-09-14T03:36:18.523788Z","iopub.status.idle":"2021-09-14T03:36:18.555315Z","shell.execute_reply.started":"2021-09-14T03:36:18.523754Z","shell.execute_reply":"2021-09-14T03:36:18.553891Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img_val = DataLoader(val_dataset,\n                        batch_size=1,\n                        num_workers=os.cpu_count(),\n                        worker_init_fn=seed_worker,\n#                                     generator=g,\n                        pin_memory=True) \nimg_val = iter(img_val)\nfor i in range(5):\n    photo_img, _ = next(img_val)\n    pred_style = gan.gen_p2s(photo_img.to(device)).cpu().detach()\n    photo_img = unnorm(photo_img).permute(0,2, 3,1)\n    pred_style = unnorm(pred_style).permute(0,2, 3,1)\n    tmp = torch.cat((photo_img, pred_style), 0)\n    if i == 0:\n        forcheck_imgs = tmp\n    else:\n        forcheck_imgs = torch.cat((forcheck_imgs,tmp),0)\nshow_imgs(forcheck_imgs,5,2)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:36:45.635443Z","iopub.execute_input":"2021-09-14T03:36:45.635855Z","iopub.status.idle":"2021-09-14T03:36:46.923697Z","shell.execute_reply.started":"2021-09-14T03:36:45.63581Z","shell.execute_reply":"2021-09-14T03:36:46.922561Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generation for submitting","metadata":{}},{"cell_type":"code","source":"class MyDataset4Out(torch.utils.data.Dataset):\n\n    def __init__(self, photo_dir, imageSize=(256,256), transform=None, normalize=True):\n        if normalize:\n            self.transform = transforms.Compose([\n                transforms.Resize(imageSize), \n                transforms.ToTensor(), \n                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n            ])\n        else:\n            self.transform = transforms.Compose([\n                transforms.Resize(imageSize), \n                transforms.ToTensor(), \n            ])\n        # Dataのパスリスト\n        self.photo_paths = [str(p) for p in Path(photo_dir).glob(\"*.jpg\")]\n\n        self.data_num = len(self.photo_paths) # ここが__len__の返り値になる\n\n    def __getitem__(self, idx):\n        p = self.photo_paths[idx]\n        photo_img = PIL.Image.open(p)\n        if self.transform:\n            out_photo = self.transform(photo_img)            \n        return out_photo\n        \n    def __len__(self):\n        return self.data_num","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:36:10.723405Z","iopub.status.idle":"2021-09-14T03:36:10.724024Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"photo_dataset = MyDataset4Out(photo_dir)\nphoto_dl = DataLoader(photo_dataset,\n                        batch_size=1,\n                        num_workers=os.cpu_count(),\n                        worker_init_fn=seed_worker,\n#                                     generator=g,\n                        pin_memory=True)\nos.makedirs('../images',exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:36:10.725179Z","iopub.status.idle":"2021-09-14T03:36:10.725772Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"t = tqdm(photo_dl, leave=False, total=photo_dl.__len__())\nfor i, photo in enumerate(t):\n    gan.gen_p2s.eval()\n    with torch.no_grad():\n        styled_photo = gan.gen_p2s(photo.to(device))\n        styled_photo = styled_photo.cpu().detach()\n    styled_photo = unnorm(styled_photo)\n    img = styled_photo[0].permute(1, 2, 0).numpy()\n    #nyumpyからpilにする\n    img_pil = PIL.Image.fromarray((img*255).astype(np.uint8))\n    img_pil.save(\"../images/\" + str(i+1) + \".jpg\")","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:36:10.726864Z","iopub.status.idle":"2021-09-14T03:36:10.727588Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:36:10.728681Z","iopub.status.idle":"2021-09-14T03:36:10.729405Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfor dirname, _, filenames in os.walk('/kaggle'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))","metadata":{"execution":{"iopub.status.busy":"2021-09-14T03:36:10.730463Z","iopub.status.idle":"2021-09-14T03:36:10.731123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}