{"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":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\nimport numpy as np\nfrom collections import OrderedDict\nimport torch.nn.functional as F\nfrom kaggle_datasets import KaggleDatasets\nimport re\nimport os\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom datetime import datetime, date, time, timedelta\nimport random\nimport timeit\nfrom abc import abstractclassmethod\nfrom venv import create\nimport tqdm\nfrom torch.utils.data import DataLoader\nimport torch\nimport matplotlib.pyplot as plt\nimport os\nfrom tqdm.auto import tqdm\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport time\nimport copy\nimport torch\nfrom torchvision.utils import make_grid\nimport matplotlib.pyplot as plt\nimport os\nfrom torchvision import transforms\nfrom torchvision.datasets import ImageFolder\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom datetime import datetime\nfrom torch.utils.data import Dataset\n","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:29.917982Z","iopub.execute_input":"2023-01-13T01:07:29.918400Z","iopub.status.idle":"2023-01-13T01:07:29.928202Z","shell.execute_reply.started":"2023-01-13T01:07:29.918369Z","shell.execute_reply":"2023-01-13T01:07:29.926857Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Bloques","metadata":{}},{"cell_type":"code","source":"\nclass ResidualBlock(nn.Module):\n    ##Bloque residual, aplicación de función identidad\n    def __init__(self, input_channels):\n        super(ResidualBlock, self).__init__()\n        self.conv1 = nn.Conv2d(input_channels, input_channels, kernel_size=3, padding=1, padding_mode='reflect')\n        self.conv2 = nn.Conv2d(input_channels, input_channels, kernel_size=3, padding=1, padding_mode='reflect')\n        self.instancenorm = nn.InstanceNorm2d(input_channels)\n        self.activation = nn.ReLU()\n\n    def forward(self, x):\n        original_x = x.clone()\n        x = self.conv1(x)\n        x = self.instancenorm(x)\n        x = self.activation(x)\n        x = self.conv2(x)\n        x = self.instancenorm(x)\n        return original_x + x ##función identidad\n\n\nclass ContractingBlock(nn.Module):\n    def __init__(self, input_channels, INorm=True, kernel_size=3, activation='reluu'):\n        super(ContractingBlock, self).__init__()\n        self.conv1 = nn.Conv2d(input_channels, input_channels * 2, kernel_size=kernel_size, padding=1, stride=2, padding_mode='reflect')\n        if activation == 'relu':\n            self.activation = nn.ReLU()\n        else:\n            self.activation = nn.LeakyReLU(0.2)\n        if INorm:\n            self.instancenorm = nn.InstanceNorm2d(input_channels * 2)\n        self.INorm = INorm\n\n    def forward(self, x):\n        x = self.conv1(x)\n        if self.INorm:\n            x = self.instancenorm(x)\n        x = self.activation(x)\n        return x\n\nclass ExpandingBlock(nn.Module):\n    def __init__(self, input_channels, INorm=True):\n        super(ExpandingBlock, self).__init__()\n        self.conv1 = nn.ConvTranspose2d(input_channels, input_channels // 2, kernel_size=3, stride=2, padding=1, output_padding=1)\n        if INorm:\n            self.instancenorm = nn.InstanceNorm2d(input_channels // 2)\n        self.INorm = INorm\n        self.activation = nn.LeakyReLU()\n\n    def forward(self, x):\n        x = self.conv1(x)\n        if self.INorm:\n            x = self.instancenorm(x)\n        x = self.activation(x)\n        return x\n\n    ##Mapping de las características de las imágenes\nclass FeatureMapBlock(nn.Module):\n    def __init__(self, input_channels, output_channels):\n        super(FeatureMapBlock, self).__init__()\n        self.conv = nn.Conv2d(input_channels, output_channels, kernel_size=7, padding=3, padding_mode='reflect')\n\n    def forward(self, x):\n        x = self.conv(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:29.933871Z","iopub.execute_input":"2023-01-13T01:07:29.934676Z","iopub.status.idle":"2023-01-13T01:07:29.952414Z","shell.execute_reply.started":"2023-01-13T01:07:29.934503Z","shell.execute_reply":"2023-01-13T01:07:29.951381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Generator","metadata":{}},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, input_channels, output_channels, hidden_channels=32):\n        super(Generator, self).__init__()\n        self.featureMB1 = FeatureMapBlock(input_channels, hidden_channels)\n        self.contractB1 = ContractingBlock(hidden_channels)\n        self.contractB2 = ContractingBlock(hidden_channels * 2)\n        res_mult = 4\n        self.residualB0 = ResidualBlock(hidden_channels * res_mult)\n        self.residualB1 = ResidualBlock(hidden_channels * res_mult)\n        self.residualB2 = ResidualBlock(hidden_channels * res_mult)\n        self.residualB3 = ResidualBlock(hidden_channels * res_mult)\n        self.residualB4 = ResidualBlock(hidden_channels * res_mult)\n        self.residualB5 = ResidualBlock(hidden_channels * res_mult)\n        self.residualB6 = ResidualBlock(hidden_channels * res_mult)\n        self.residualB7 = ResidualBlock(hidden_channels * res_mult)\n        self.residualB8 = ResidualBlock(hidden_channels * res_mult)\n\n        self.expandB1 = ExpandingBlock(hidden_channels * 4)\n        self.expandB2 = ExpandingBlock(hidden_channels * 2)\n        self.featureMB2 = FeatureMapBlock(hidden_channels, output_channels)\n        self.tanh = nn.Tanh()\n\n    def forward(self, x):\n        x0 = self.featureMB1(x)\n        x1 = self.contractB1(x0)\n        x2 = self.contractB2(x1)\n        x3 = self.residualB0(x2)\n        x4 = self.residualB1(x3)\n        x5 = self.residualB2(x4)\n        x6 = self.residualB3(x5)\n        x7 = self.residualB4(x6)\n        x8 = self.residualB5(x7)\n        x9 = self.residualB6(x8)\n        x10 = self.residualB7(x9)\n        x11 = self.residualB8(x10)\n        x12 = self.expandB1(x11)\n        x13 = self.expandB2(x12)\n        xn = self.featureMB2(x13)\n        return self.tanh(xn)","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:29.954931Z","iopub.execute_input":"2023-01-13T01:07:29.955330Z","iopub.status.idle":"2023-01-13T01:07:29.972840Z","shell.execute_reply.started":"2023-01-13T01:07:29.955293Z","shell.execute_reply":"2023-01-13T01:07:29.971680Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Discriminator","metadata":{}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, input_channels, hidden_channels=64):\n        super(Discriminator, self).__init__()\n        self.featureMB = FeatureMapBlock(input_channels, hidden_channels)\n        self.contractB1 = ContractingBlock(hidden_channels, INorm=False, kernel_size=4, activation='lrelu')\n        self.contractB2 = ContractingBlock(hidden_channels * 2, kernel_size=4, activation='lrelu')\n        self.contractB3 = ContractingBlock(hidden_channels * 4, kernel_size=4, activation='lrelu')\n        self.outputLayer = nn.Conv2d(hidden_channels * 8, 1, kernel_size=1)\n\n    def forward(self, x):\n        x0 = self.featureMB(x)\n        x1 = self.contractB1(x0)\n        x2 = self.contractB2(x1)\n        x3 = self.contractB3(x2)\n        xn = self.outputLayer(x3)\n        return xn\n","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:29.975826Z","iopub.execute_input":"2023-01-13T01:07:29.976486Z","iopub.status.idle":"2023-01-13T01:07:29.988370Z","shell.execute_reply.started":"2023-01-13T01:07:29.976460Z","shell.execute_reply":"2023-01-13T01:07:29.987287Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Training","metadata":{}},{"cell_type":"code","source":"class CycleTrainer():\n    def __init__(self, dataloader1, dataloader2, genAB, genBA, genOptim, \n                 discA, discAOptim, discB, discBOptim,\n                 advCriterion,\n                 reconCriterion, logStep, logDir, targetShape, device='cuda', \n                 checksave=False, saveStep=None,\n                 load=False,\n                 loadDir=None, genDiscLoad=None, timeSteps=False, timeEpochs=False):\n\n        self.device = device\n        self.advCriterion = advCriterion\n        self.reconCriterion = reconCriterion\n        self.logStep = logStep\n        self.logDir = logDir\n        self.checksave = checksave\n        self.timeEpochs = timeEpochs\n        self.timeSteps = timeSteps\n        self.targetShape = targetShape\n        self.genAB = genAB\n        self.genBA = genBA\n        self.genOptim = genOptim\n        self.discA = discA\n        self.discAOptim = discAOptim\n        self.discB = discB\n        self.discBOptim = discBOptim\n        self.genLoss = []\n        self.disLoss = []\n        self.dataloader1 = dataloader1\n        self.dataloader2 = dataloader2\n\n\n        if checksave:\n            self.saveStep = saveStep\n\n        self.load = load\n\n        if load:\n            self.loadDir = loadDir\n            self.genDiscLoad = genDiscLoad\n\n        self.act = 0\n        self.iter = 0\n\n        if timeSteps:\n            self.stepTimes = []\n            self.numSteps = []\n        if timeEpochs:\n            self.epochTimes = []\n            self.numEpochs = []\n        \n        if load:\n            self.load_checkpoint()\n        else:\n            genAB = genAB.apply(self.weightsInit)\n            genBA = genBA.apply(self.weightsInit)\n            discA = discA.apply(self.weightsInit)\n            discB = discB.apply(self.weightsInit)\n\n    def getDiscLoss(self, real, fake, discrim, advCriterion):\n        ##Obtenemos resultados del discriminador y calculamos su pérdida\n        ##Para imágenes falsas e imágenes reales\n        discFakeresult = discrim(fake.detach()) \n        discFakeLoss = advCriterion(discFakeresult, \n                                    torch.zeros_like(discFakeresult))\n        discRealresult = discrim(real)\n        discRealLoss = advCriterion(discRealresult, \n                                         torch.ones_like(discRealresult))\n        disc_loss = (discFakeLoss + discRealLoss) / 2\n        return disc_loss\n\n    def getGenAdversarialLoss(self, real, discrim, genA, advCriterion):\n        fake = genA(real)\n        discFakePred= discrim(fake)\n        adversarialLoss = advCriterion(discFakePred, \n                                       torch.ones_like(discFakePred))\n        return adversarialLoss, fake\n\n    def getIdentityLoss(self, real, genB, criterionI):\n        identityResGen = genB(real)\n        identityLoss = criterionI(identityResGen, real)\n        return identityLoss\n\n    def getCycleConsistencyLoss(self, real, fake, genB, cycleCriterion):\n        fakeRes = genB(fake)\n        cycleLoss = cycleCriterion(fakeRes, real)\n        return cycleLoss\n\n    def getGenLoss(self, realA, realB, genAB, genBA, discA, discB, advCriterion, identityCriterion,\n                     cycleCriterion, lambdaIdentity=0.1, lambdaCycle=10):\n        \n        lossAdvBA, fakeA = self.getGenAdversarialLoss(realB, discA, genBA, advCriterion)\n        lossAdvAB, fakeB = self.getGenAdversarialLoss(realA, discB, genAB, advCriterion)\n        genAdversarialLoss = lossAdvBA + lossAdvAB\n\n        \n        lossIdentityA = self.getIdentityLoss(realA, genBA, identityCriterion)\n        lossIdentityB  = self.getIdentityLoss(realB, genAB, identityCriterion)\n        genIdentityLoss = lossIdentityA + lossIdentityB\n\n        \n        lossCycleBA = self.getCycleConsistencyLoss(realA, fakeB, genBA, cycleCriterion)\n        lossCycleAB = self.getCycleConsistencyLoss(realB, fakeA, genAB, cycleCriterion)\n        genCycleLoss = lossCycleBA + lossCycleAB\n\n        # Total\n        genLoss = lambdaIdentity * genIdentityLoss + lambdaCycle * genCycleLoss + genAdversarialLoss\n        \n        return genLoss, fakeA, fakeB\n\n    def weightsInit(self, m):\n        if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n            torch.nn.init.normal_(m.weight, 0.0, 0.02)\n        if isinstance(m, nn.BatchNorm2d):\n            torch.nn.init.normal_(m.weight, 0.0, 0.02)\n            torch.nn.init.constant_(m.bias, 0)\n\n    def loadCheckpoint(self):\n        os.chdir(self.loadDir)\n\n        pre_dict = torch.load(self.gen_disc_load)\n\n        self.genAB.load_state_dict(pre_dict['genAB'])\n        self.genBA.load_state_dict(pre_dict['genBA'])\n        self.genOptim.load_state_dict(pre_dict['genOptim'])\n\n        self.discA.load_state_dict(pre_dict['discA'])\n        self.discAOptim.load_state_dict(pre_dict['discAOptim'])\n        self.discB.load_state_dict(pre_dict['discB'])\n        self.discBOptim.load_state_dict(pre_dict['discBOptim'])\n\n        self.act = pre_dict['epoch']\n\n        self.disLoss = pre_dict['discLoss']\n\n        self.genLoss = pre_dict['genLoss']\n\n        os.chdir('..')\n\n    def saveCheckpoint(self, epoch):\n        os.chdir(self.logDir)\n\n        torch.save({\n            'genAB': self.genAB.state_dict(),\n            'genBA': self.genBA.state_dict(),\n            'genOptim': self.genOptim.state_dict(),\n            'discA': self.discA.state_dict(),\n            'discAOptim': self.discAOptim.state_dict(),\n            'discB': self.discB.state_dict(),\n            'discBOptim': self.discBOptim.state_dict(),\n            'epoch': epoch,\n            'genLoss': self.genLoss,\n            'discLoss': self.disLoss\n        }, 'cycleGAN_' + str(epoch) + '.pth')\n\n        os.chdir('..')\n\n    def enableTraining(self, model, flag):\n        for p in model.parameters():\n            p.requires_grad = flag\n            \n    def trainGen(self, realA, realB):\n        self.enableTraining(self.genBA, True)\n        self.enableTraining(self.genAB, True)\n        self.enableTraining(self.discA, False)\n        self.enableTraining(self.discB, False)\n\n        ### GET GEN LOSS \n        self.genOptim.zero_grad()\n        genLoss, fakeA, fakeB = self.getGenLoss(\n            realA, realB, self.genAB, self.genBA, self.discA, self.discB, self.advCriterion,\n            self.reconCriterion, self.reconCriterion\n        )\n        genLoss.backward()  # Gradients\n        self.genOptim.step()  # Optimizer\n\n        return (genLoss, fakeA, fakeB)\n\n    def trainDisc(self, realA, realB):\n        self.enableTraining(self.genBA, False)\n        self.enableTraining(self.genAB, False)\n        self.enableTraining(self.discA, True)\n        self.enableTraining(self.discB, True)\n\n        ### GET LOSS DISC A \n        self.discAOptim.zero_grad()  \n        with torch.no_grad():\n            fakeA = self.genBA(realB)\n        discALoss = self.getDiscLoss(realA, fakeA, self.discA, self.advCriterion)\n        discALoss.backward(retain_graph=True)  # Gradients\n        self.discAOptim.step()  # Optimizer\n\n        ### GET LOSS DISC B\n        self.discBOptim.zero_grad()  \n        with torch.no_grad():\n            fakeB = self.genAB(realA)\n        discBLoss = self.getDiscLoss(realB, fakeB, self.discB, self.advCriterion)\n        discBLoss.backward(retain_graph=True)  # Gradients\n        self.discBOptim.step()  # Optimizer\n\n        return discALoss, discBLoss\n\n    def trainForEpochs(self, n_epochs):\n        self.initialTime = time.time()\n        self.meanGeneratorLoss = 0\n        self.meanDiscriminatorLoss = 0\n        for e in range(0, n_epochs):\n            self.epoch(e + self.act)\n            self.act += 1\n\n    def epoch(self, ep):\n        it = 0\n        for (realA, realB) in tqdm(zip(self.dataloader1, self.dataloader2)):\n            realA = nn.functional.interpolate(realA, size=self.targetShape)\n            realB = nn.functional.interpolate(realB, size=self.targetShape)\n            realA = realA.to(self.device)\n            realB = realB.to(self.device)\n\n            #Train GEN and DIS\n            generatorLoss, fakeA, fakeB = self.trainGen(realA, realB)\n            discALoss, discBLoss = self.trainDisc(realA, realB)\n\n            # update Dis loss\n            self.meanDiscriminatorLoss += discALoss.item() / self.logStep\n            \n            # update Gen loss\n            self.meanGeneratorLoss += generatorLoss.item() / self.logStep\n           \n            self.genLoss.append(generatorLoss.item())\n           \n            self.disLoss.append(discALoss.item())\n            \n           \n\n            ### Visualization\n            if self.iter % self.logStep == 0 and self.iter > 0:\n                #self.plotLosses()\n                self.saveResults(torch.cat([realA, fakeA]), torch.cat([realB, fakeA]))\n               \n                #self.saveResults(fakeB, fakeA)\n\n\n                #if self.time_steps:\n                #    self.plot_step_time()\n\n            #if self.checksave and self.iter % self.saveStep == 0:\n            #    self.saveCheckpoint(ep)\n                \n            self.iter = self.iter + 1\n            it = it + 1\n\n\n    def saveResults(self, real, generated):\n        tensor_as_image(generated, self.iter, \"fake\", self.logDir, save=False, show=True)\n        tensor_as_image(real, self.iter, \"real\", self.logDir, save=False, show=True)\n\n    def enableTraining(self, model, flag):\n        for p in model.parameters():\n            p.requires_grad = flag\n\n    def checkTrainingParams(self, model, flag):\n        for p in model.parameters():\n            assert (p.requires_grad == flag)\n\n    def plotEpochTime(self):\n        epochTime = (time.time() - self.initial__time)\n        self.epochTimes.append(epoch_time)\n        self.numEpochs.append(self.act)\n\n        title = str(self.epochTimes[-1]) + \" seconds taken\"\n        print(title)\n\n        plt.plot(self.numepochs, self.epochTimes, label=\"Epoch Time\")\n\n        plt.title(title)\n        plt.legend()\n\n        os.chdir(self.logDir)\n        plt.savefig(datetime.now().strftime(\"%d-%m\") + \" iter \" + str(self.act) + \" epoch_times\" + '.pdf')\n        os.chdir('..')\n\n        plt.clf()\n\n    def plotStepTime(self):\n\n        stepTime = (time.time() - self.initialTime)\n        self.stepTimes.append(stepTime)\n        self.numSteps.append(self.iter)\n\n        title = str(self.stepTimes[-1]) + \" seconds to take \" + str(self.iter) + \" steps\"\n        print(title)\n\n        # Time visualization\n\n        plt.plot(self.numSteps, self.stepTimes, label=\"Step Time\")\n\n        plt.title(title)\n        plt.legend()\n\n        os.chdir(self.logDir)\n        plt.savefig(datetime.now().strftime(\"%d-%m\") + \" iter \" + str(self.iter) + \" step_times\" + '.pdf')\n        os.chdir('..')\n\n        plt.clf()\n\n    def plotLosses(self):\n\n        self.meanGeneratorLoss = 0\n        self.meanDiscriminatorLoss = 0\n        stepBins = 20\n        xAxis = sorted([i * stepBins for i in range(len(self.genLoss) // stepBins)] * stepBins)\n        numExamples = (len(self.genLoss) // stepBins) * stepBins\n        title = \" GEN LOSS:\" + str(sum(self.genLoss[:numExamples]) / len(self.genLoss[:numExamples]))\n        print(title)\n        title = \" DIS LOSS:\" + str(sum(self.dis_loss[:num_examples]) / len(self.disLoss[:numExamples]))\n        print(title)\n        plt.plot(\n            range(numExamples // stepBins),\n            torch.Tensor(self.genLoss[:numExamples]).view(-1, stepBins).mean(1),\n            label=\"Generator Loss\"\n        )\n        plt.plot(\n            range(numExamples // stepBins),\n            torch.Tensor(self.dis_loss[:numExamples]).view(-1, stepBins).mean(1),\n            label=\"Discriminator Loss\"\n        )\n        plt.legend()\n\n\n        os.chdir(self.logDir)\n        plt.savefig(datetime.now().strftime(\"%d-%m\") + \" iter \" + str(self.iter) + '.pdf')\n        os.chdir('..')\n        plt.clf()\n\n    def generateSamples(self, nSamples=10000):\n        i = 0\n        for (realA, realB) in tqdm(zip(self.dataloader1, self.dataloader2)):\n            i += 1\n            realA = nn.functional.interpolate(real_A, size=self.target_shape)\n            realB = nn.functional.interpolate(real_B, size=self.target_shape)\n            realA = real_A.to(self.device)\n            realB = real_B.to(self.device)\n            genLoss, fakeA, fakeB = self.getGenLoss(\n                realA, realB, self.genAB, self.genBA, self.discA, self.discB, self.advCriterion,\n                self.recon_criterion, self.recon_criterion\n            )\n            self.saveResults(fakeB, fakeA)\n            #self.save_results(torch.cat([real_A, real_B]), torch.cat([fake_B, fake_A]))\n            #ImageFunctions.tensor_as_image(fake_A, self.iter + random.randint(0, 11119), \"newData\", self.log_dir,\n            #                               save=True, show=False)\n            ##ImageFunctions.tensor_as_image(fake_B, self.iter + random.randint(0, 11119), \"newData\", self.log_dir,\n            ##                               save=True, show=False)\n\n            ##self.save_results(torch.cat([real_A, real_B]), torch.cat([fake_B, fake_A]))\n\n            if i > n_samples:\n                break","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:30.030438Z","iopub.execute_input":"2023-01-13T01:07:30.031074Z","iopub.status.idle":"2023-01-13T01:07:30.086200Z","shell.execute_reply.started":"2023-01-13T01:07:30.031039Z","shell.execute_reply":"2023-01-13T01:07:30.085210Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"data","metadata":{}},{"cell_type":"code","source":"datafolder = KaggleDatasets().get_gcs_path()\n","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:30.088274Z","iopub.execute_input":"2023-01-13T01:07:30.088699Z","iopub.status.idle":"2023-01-13T01:07:30.556472Z","shell.execute_reply.started":"2023-01-13T01:07:30.088654Z","shell.execute_reply":"2023-01-13T01:07:30.555440Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import Dataset\n\nclass crearDataset(Dataset):\n    def __setup_files(self):\n        fileNames = os.listdir(self.dataDir)\n        files = []  # initialize to empty list\n\n        for fileName in fileNames:\n            pFile = os.path.join(self.dataDir, fileName)\n            if os.path.isfile(pFile):\n                files.append(pFile)\n\n        return files\n\n    def __init__(self, dataDir, transform=None):\n        self.dataDir = dataDir\n        self.transform = transform ##Tranformaciones \n        self.files = self.__setup_files()\n\n    def __len__(self):\n        return len(self.files)\n\n    def __getitem__(self, idx):\n        from PIL import Image\n\n        imgFile = self.files[idx]\n\n        if imgFile[-4:] == \".npy\":\n            # files are in .npy format\n            img = np.load(imgFile)\n            img = Image.fromarray(img.squeeze(0).transpose(1, 2, 0))\n\n        else:\n            # read the image:\n            img = Image.open(self.files[idx]).convert('RGB')\n\n        # apply the transforms on the image\n        if self.transform is not None:\n            img = self.transform(img)\n\n        if img.shape[0] >= 4:\n            img = img[:3, :, :]\n\n        # return the image:\n        return img\n\nfrom torchvision.transforms import ToTensor, Normalize, Compose, Resize, RandomHorizontalFlip\n\ndef makeDataset(resolution, imgDir):\n\n    transforms = Compose([RandomHorizontalFlip(),Resize(resolution),\n                          ToTensor(),Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])\n\n    return crearDataset(imgDir, transform=transforms)","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:30.558310Z","iopub.execute_input":"2023-01-13T01:07:30.559011Z","iopub.status.idle":"2023-01-13T01:07:30.579492Z","shell.execute_reply.started":"2023-01-13T01:07:30.558974Z","shell.execute_reply":"2023-01-13T01:07:30.578470Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import random\n\ndef tensor_as_image(img_tensor, iter, tag=None, dir=None, num_images=25, size=(3, 64, 64), nrow=5, save=False,\n                    show=True):\n    img_tensor = (img_tensor + 1) / 2\n    image_unflat = img_tensor.detach().cpu()\n    image_grid = make_grid(image_unflat[:num_images], nrow=nrow)\n    plt.imshow(image_grid.permute(1, 2, 0).squeeze())\n    plt.axis('off')\n    if save:\n        plt.savefig(dir + datetime.now().strftime(\"%d-%m\") + \" iter \" + str(iter) + tag + '_'+ str(random.randint(0, 11119)) +'.jpg', bbox_inches='tight')\n    if show:\n        plt.show()\n\n    plt.clf()","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:30.583885Z","iopub.execute_input":"2023-01-13T01:07:30.584616Z","iopub.status.idle":"2023-01-13T01:07:30.598310Z","shell.execute_reply.started":"2023-01-13T01:07:30.584581Z","shell.execute_reply":"2023-01-13T01:07:30.597233Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"nEpochs = 50\nbatchSize = 1\nlr = 0.00005\ntargetShape = 256\ndevice = 'cuda'\n\n# Modulos\ngenAB = Generator(3, 3).to(device)\ngenBA = Generator(3, 3).to(device)\ngenOptim = torch.optim.Adam(list(genAB.parameters()) + list(genBA.parameters()), lr=lr, betas=(0.5, 0.999))\ndiscA = Discriminator(3).to(device)\ndiscAOptim = torch.optim.Adam(discA.parameters(), lr=lr, betas=(0.5, 0.999))\ndiscB = Discriminator(3).to(device)\ndiscBOptim = torch.optim.Adam(discB.parameters(), lr=lr, betas=(0.5, 0.999))\n\n\nadvCriterion = nn.MSELoss() \nreconCriterion = nn.L1Loss() \n\n# Carpeta para resultados\n\nimgDir1=\"../input/gan-getting-started/monet_jpg\"\nimgDir2=\"../input/gan-getting-started/photo_jpg\"\n\n# Dataset\ndataset1 = makeDataset(resolution=256,imgDir=imgDir1)\ndataset2 = makeDataset(resolution=256, imgDir=imgDir2)\n\ndataLoader1 = DataLoader(dataset1, batch_size=batchSize, shuffle=True)\ndataLoader2 = DataLoader(dataset2, batch_size=batchSize, shuffle=True)\n\ncriterion = nn.BCEWithLogitsLoss()\ndisplayStep = 200\ncheckpointStep = 200\nload=False\nchecksave = False\ntrainingDir = 'CycleTraining'\nloadDir = 'models'\ngenDiscLoad = 'cycleGAN_256.pth'\n\n","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:30.604690Z","iopub.execute_input":"2023-01-13T01:07:30.605706Z","iopub.status.idle":"2023-01-13T01:07:33.895793Z","shell.execute_reply.started":"2023-01-13T01:07:30.605669Z","shell.execute_reply":"2023-01-13T01:07:33.894784Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\ntrainer = CycleTrainer(dataLoader1, dataLoader2, \n                                 genAB, genBA, genOptim, \n                                 discA, discAOptim, discB, \n                                 discBOptim, advCriterion,\n                                 reconCriterion, displayStep, \n                                 trainingDir, targetShape, 'cuda', \n                                 checksave, checkpointStep, load, loadDir, \n                                 genDiscLoad, timeSteps = True, timeEpochs = True)\n","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:33.900518Z","iopub.execute_input":"2023-01-13T01:07:33.903394Z","iopub.status.idle":"2023-01-13T01:07:33.916799Z","shell.execute_reply.started":"2023-01-13T01:07:33.903357Z","shell.execute_reply":"2023-01-13T01:07:33.915866Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.trainForEpochs(nEpochs)","metadata":{"execution":{"iopub.status.busy":"2023-01-13T01:07:33.922202Z","iopub.execute_input":"2023-01-13T01:07:33.925269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import PIL\n! mkdir /kaggle/working/images","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ls /kaggle/working/","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"iter = 1\nimport os\nos.chdir(\"/kaggle/working/\")\n\nimport zipfile\nwith zipfile.ZipFile(\"images.zip\", mode=\"w\") as archive:\n\n    for img in dataLoader2:         \n        num_images =1 \n        nrow=5\n        realB = nn.functional.interpolate(img, size=256)\n        realB = realB.to('cuda')\n        generated = genBA(realB)\n        img_tensor = generated\n        img_tensor = (img_tensor + 1) / 2\n        image_unflat = img_tensor.detach().cpu()\n        image_grid = make_grid(image_unflat[:num_images], nrow=nrow)\n        plt.imshow(image_grid.permute(1, 2, 0).squeeze())\n        plt.axis('off')\n        plt.savefig( \"iter\" + str(iter) + '_'+'.jpg', bbox_inches='tight')\n        archive.write(\"iter\" + str(iter) + '_'+'.jpg')\n        os.remove(\"iter\" + str(iter) + '_'+'.jpg')\n        plt.clf()\n        iter += 1\n  \n    ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}