\n\n![](\"https://junyanz.github.io/CycleGAN/images/teaser_high_res.jpg\")
\n\n[Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://junyanz.github.io/CycleGAN/) \n [Jun-Yan Zhu](https://people.eecs.berkeley.edu/~junyanz/)\\*, [Taesung Park](https://taesung.me/)\\*, [Phillip Isola](http://web.mit.edu/phillipi/), [Alexei A. Efros](https://people.eecs.berkeley.edu/~efros/) \n Berkeley AI Research Lab, UC Berkeley \n In ICCV 2017. (* equal contributions) \n\nThis package includes CycleGAN, [pix2pix](https://github.com/phillipi/pix2pix), as well as other methods like [BiGAN](https://arxiv.org/abs/1605.09782)/[ALI](https://ishmaelbelghazi.github.io/ALI/) and Apple's paper [S+U learning](https://arxiv.org/pdf/1612.07828.pdf). \nThe code was written by [Jun-Yan Zhu](https://github.com/junyanz) and [Taesung Park](https://github.com/taesung). \n**Update**: Please check out [PyTorch](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) implementation for CycleGAN and pix2pix.\nThe PyTorch version is under active development and can produce results comparable or better than this Torch version.\n\n## Other implementations:\n![](\"https://junyanz.github.io/CycleGAN/images/painting2photo.jpg\")
\n\n### Collection Style Transfer\n![](\"https://junyanz.github.io/CycleGAN/images/photo2painting.jpg\")
\n\n### Object Transfiguration\n![](\"https://junyanz.github.io/CycleGAN/images/objects.jpg\")
\n\n### Season Transfer\n![](\"https://junyanz.github.io/CycleGAN/images/season.jpg\")
\n\n### Photo Enhancement: Narrow depth of field\n![](\"https://junyanz.github.io/CycleGAN/images/photo_enhancement.jpg\")
\n\n\n\n## Prerequisites\n- Linux or OSX\n- NVIDIA GPU + CUDA CuDNN (CPU mode and CUDA without CuDNN may work with minimal modification, but untested)\n- For MAC users, you need the Linux/GNU commands `gfind` and `gwc`, which can be installed with `brew install findutils coreutils`.\n\n## Getting Started\n### Installation\n- Install torch and dependencies from https://github.com/torch/distro\n- Install torch packages `nngraph`, `class`, `display`\n```bash\nluarocks install nngraph\nluarocks install class\nluarocks install https://raw.githubusercontent.com/szym/display/master/display-scm-0.rockspec\n```\n- Clone this repo:\n```bash\ngit clone https://github.com/junyanz/CycleGAN\ncd CycleGAN\n```\n\n### Apply a Pre-trained Model\n- Download the test photos (taken by [Alexei Efros](https://www.flickr.com/photos/aaefros)):\n```\nbash ./datasets/download_dataset.sh ae_photos\n```\n- Download the pre-trained model `style_cezanne` (For CPU model, use `style_cezanne_cpu`):\n```\nbash ./pretrained_models/download_model.sh style_cezanne\n```\n- Now, let's generate Paul Cézanne style images:\n```\nDATA_ROOT=./datasets/ae_photos name=style_cezanne_pretrained model=one_direction_test phase=test loadSize=256 fineSize=256 resize_or_crop=\"scale_width\" th test.lua\n```\nThe test results will be saved to `./results/style_cezanne_pretrained/latest_test/index.html`. \nPlease refer to [Model Zoo](#model-zoo) for more pre-trained models.\n`./examples/test_vangogh_style_on_ae_photos.sh` is an example script that downloads the pretrained Van Gogh style network and runs it on Efros's photos.\n\n### Train\n- Download a dataset (e.g. zebra and horse images from ImageNet):\n```bash\nbash ./datasets/download_dataset.sh horse2zebra\n```\n- Train a model:\n```bash\nDATA_ROOT=./datasets/horse2zebra name=horse2zebra_model th train.lua\n```\n- (CPU only) The same training command without using a GPU or CUDNN. Setting the environment variables ```gpu=0 cudnn=0``` forces CPU only\n```bash\nDATA_ROOT=./datasets/horse2zebra name=horse2zebra_model gpu=0 cudnn=0 th train.lua\n```\n- (Optionally) start the display server to view results as the model trains. (See [Display UI](#display-ui) for more details):\n```bash\nth -ldisplay.start 8000 0.0.0.0\n```\n\n### Test\n- Finally, test the model:\n```bash\nDATA_ROOT=./datasets/horse2zebra name=horse2zebra_model phase=test th test.lua\n```\nThe test results will be saved to an HTML file here: `./results/horse2zebra_model/latest_test/index.html`.\n\n\n## Model Zoo\nDownload the pre-trained models with the following script. The model will be saved to `./checkpoints/model_name/latest_net_G.t7`.\n```bash\nbash ./pretrained_models/download_model.sh model_name\n```\n- `orange2apple` (orange -> apple) and `apple2orange`: trained on ImageNet categories `apple` and `orange`.\n- `horse2zebra` (horse -> zebra) and `zebra2horse` (zebra -> horse): trained on ImageNet categories `horse` and `zebra`.\n- `style_monet` (landscape photo -> Monet painting style), `style_vangogh` (landscape photo -> Van Gogh painting style), `style_ukiyoe` (landscape photo -> Ukiyo-e painting style), `style_cezanne` (landscape photo -> Cezanne painting style): trained on paintings and Flickr landscape photos.\n- `monet2photo` (Monet paintings -> real landscape): trained on paintings and Flickr landscape photographs.\n- `cityscapes_photo2label` (street scene -> label) and `cityscapes_label2photo` (label -> street scene): trained on the Cityscapes dataset.\n- `map2sat` (map -> aerial photo) and `sat2map` (aerial photo -> map): trained on Google maps.\n- `iphone2dslr_flower` (iPhone photos of flowers -> DSLR photos of flowers): trained on Flickr photos.\n\nCPU models can be downloaded using:\n```bash\nbash pretrained_models/download_model.sh ![](\"https://junyanz.github.io/CycleGAN/images/failure_putin.jpg\")
\n\nOur model does not work well when the test image is rather different from the images on which the model is trained, as is the case in the figure to the left (we trained on horses and zebras without riders, but test here one a horse with a rider). See additional typical failure cases [here](https://junyanz.github.io/CycleGAN/images/failures.jpg). On translation tasks that involve color and texture changes, like many of those reported above, the method often succeeds. We have also explored tasks that require geometric changes, with little success. For example, on the task of `dog<->cat` transfiguration, the learned translation degenerates into making minimal changes to the input. We also observe a lingering gap between the results achievable with paired training data and those achieved by our unpaired method. In some cases, this gap may be very hard -- or even impossible,-- to close: for example, our method sometimes permutes the labels for tree and building in the output of the cityscapes photos->labels task.\n\n\n\n## Citation\nIf you use this code for your research, please cite our [paper](https://junyanz.github.io/CycleGAN/):\n\n```\n@inproceedings{CycleGAN2017,\n title={Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networkss},\n author={Zhu, Jun-Yan and Park, Taesung and Isola, Phillip and Efros, Alexei A},\n booktitle={Computer Vision (ICCV), 2017 IEEE International Conference on},\n year={2017}\n}\n\n```\n\n\n## Related Projects:\n**[contrastive-unpaired-translation](https://github.com/taesungp/contrastive-unpaired-translation) (CUT)**\n**[pix2pix-Torch](https://github.com/phillipi/pix2pix) | [pix2pixHD](https://github.com/NVIDIA/pix2pixHD) |\n[BicycleGAN](https://github.com/junyanz/BicycleGAN) | [vid2vid](https://tcwang0509.github.io/vid2vid/) | [SPADE/GauGAN](https://github.com/NVlabs/SPADE)**
\n**[iGAN](https://github.com/junyanz/iGAN) | [GAN Dissection](https://github.com/CSAILVision/GANDissect) | [GAN Paint](http://ganpaint.io/)**\n\n## Cat Paper Collection\nIf you love cats, and love reading cool graphics, vision, and ML papers, please check out the Cat Paper [Collection](https://github.com/junyanz/CatPapers). \n\n\n## Acknowledgments\nCode borrows from [pix2pix](https://github.com/phillipi/pix2pix) and [DCGAN](https://github.com/soumith/dcgan.torch). The data loader is modified from [DCGAN](https://github.com/soumith/dcgan.torch) and [Context-Encoder](https://github.com/pathak22/context-encoder). The generative network is adopted from [neural-style](https://github.com/jcjohnson/neural-style) with [Instance Normalization](https://github.com/DmitryUlyanov/texture_nets/blob/master/InstanceNormalization.lua).\n" }, { "path": "data/aligned_data_loader.lua", "content": "--------------------------------------------------------------------------------\n-- Subclass of BaseDataLoader that provides data from two datasets.\n-- The samples from the datasets are aligned\n-- The datasets are of the same size\n--------------------------------------------------------------------------------\nrequire 'data.base_data_loader'\n\nlocal class = require 'class'\ndata_util = paths.dofile('data_util.lua')\n\nAlignedDataLoader = class('AlignedDataLoader', 'BaseDataLoader')\n\nfunction AlignedDataLoader:__init(conf)\n BaseDataLoader.__init(self, conf)\n conf = conf or {}\nend\n\nfunction AlignedDataLoader:name()\n return 'AlignedDataLoader'\nend\n\nfunction AlignedDataLoader:Initialize(opt)\n opt.align_data = 1\n self.idx_A = {1, opt.input_nc}\n self.idx_B = {opt.input_nc+1, opt.input_nc+opt.output_nc}\n local nc = 3--opt.input_nc + opt.output_nc\n self.data = data_util.load_dataset('', opt, nc)\nend\n\n-- actually fetches the data\n-- |return|: a table of two tables, each corresponding to\n-- the batch for dataset A and dataset B\nfunction AlignedDataLoader:LoadBatchForAllDatasets()\n local batch_data, path = self.data:getBatch()\n local batchA = batch_data[{ {}, self.idx_A, {}, {} }]\n local batchB = batch_data[{ {}, self.idx_B, {}, {} }]\n\n return batchA, batchB, path, path\nend\n\n-- returns the size of each dataset\nfunction AlignedDataLoader:size(dataset)\n return self.data:size()\nend\n" }, { "path": "data/base_data_loader.lua", "content": "--------------------------------------------------------------------------------\n-- Base Class for Providing Data\n--------------------------------------------------------------------------------\n\nlocal class = require 'class'\nrequire 'torch'\n\nBaseDataLoader = class('BaseDataLoader')\n\nfunction BaseDataLoader:__init(conf)\n conf = conf or {} \n self.data_tm = torch.Timer()\nend\n\nfunction BaseDataLoader:name()\n return 'BaseDataLoader'\nend\n\nfunction BaseDataLoader:Initialize(opt) \nend\n\n-- actually fetches the data\n-- |return|: a table of two tables, each corresponding to \n-- the batch for dataset A and dataset B\nfunction BaseDataLoader:LoadBatchForAllDatasets()\n return {},{},{},{}\nend\n\n-- returns the next batch\n-- a wrapper of getBatch(), which is meant to be overriden by subclasses\n-- |return|: a table of two tables, each corresponding to \n-- the batch for dataset A and dataset B\nfunction BaseDataLoader:GetNextBatch()\n self.data_tm:reset()\n self.data_tm:resume()\n local dataA, dataB, pathA, pathB = self:LoadBatchForAllDatasets()\n self.data_tm:stop()\n return dataA, dataB, pathA, pathB\nend\n\nfunction BaseDataLoader:time_elapsed_to_fetch_data()\n return self.data_tm:time().real\nend\n\n-- returns the size of each dataset\nfunction BaseDataLoader:size(dataset)\n return 0\nend\n\n\n\n\n\n" }, { "path": "data/data.lua", "content": "--[[\n This data loader is a modified version of the one from dcgan.torch\n (see https://github.com/soumith/dcgan.torch/blob/master/data/data.lua).\n\n Copyright (c) 2016, Deepak Pathak [See LICENSE file for details]\n]]--\n\nlocal Threads = require 'threads'\nThreads.serialization('threads.sharedserialize')\n\nlocal data = {}\n\nlocal result = {}\nlocal unpack = unpack and unpack or table.unpack\n\nfunction data.new(n, opt_)\n opt_ = opt_ or {}\n local self = {}\n for k,v in pairs(data) do\n self[k] = v\n end\n\n local donkey_file = 'donkey_folder.lua'\n-- print('n..' .. n)\n if n > 0 then\n local options = opt_\n self.threads = Threads(n,\n function() require 'torch' end,\n function(idx)\n opt = options\n tid = idx\n local seed = (opt.manualSeed and opt.manualSeed or 0) + idx\n torch.manualSeed(seed)\n torch.setnumthreads(1)\n print(string.format('Starting donkey with id: %d seed: %d', tid, seed))\n assert(options, 'options not found')\n assert(opt, 'opt not given')\n print(opt)\n paths.dofile(donkey_file)\n end\n\n )\n else\n if donkey_file then paths.dofile(donkey_file) end\n-- print('empty threads')\n self.threads = {}\n function self.threads:addjob(f1, f2) f2(f1()) end\n function self.threads:dojob() end\n function self.threads:synchronize() end\n end\n\n local nSamples = 0\n self.threads:addjob(function() return trainLoader:size() end,\n function(c) nSamples = c end)\n self.threads:synchronize()\n self._size = nSamples\n\n for i = 1, n do\n self.threads:addjob(self._getFromThreads,\n self._pushResult)\n end\n-- print(self.threads)\n return self\nend\n\nfunction data._getFromThreads()\n assert(opt.batchSize, 'opt.batchSize not found')\n return trainLoader:sample(opt.batchSize)\nend\n\nfunction data._pushResult(...)\n local res = {...}\n if res == nil then\n self.threads:synchronize()\n end\n result[1] = res\nend\n\n\n\nfunction data:getBatch()\n -- queue another job\n self.threads:addjob(self._getFromThreads, self._pushResult)\n self.threads:dojob()\n local res = result[1]\n\n img_data = res[1]\n img_paths = res[3]\n\n result[1] = nil\n if torch.type(img_data) == 'table' then\n img_data = unpack(img_data)\n end\n\n\n return img_data, img_paths\nend\n\nfunction data:size()\n return self._size\nend\n\nreturn data\n" }, { "path": "data/data_util.lua", "content": "local data_util = {}\n\nrequire 'torch'\n-- options = require '../options.lua'\n-- load dataset from the file system\n-- |name|: name of the dataset. It's currently either 'A' or 'B'\nfunction data_util.load_dataset(name, opt, nc)\n local tensortype = torch.getdefaulttensortype()\n torch.setdefaulttensortype('torch.FloatTensor')\n\n local new_opt = options.clone(opt)\n new_opt.manualSeed = torch.random(1, 10000) -- fix seed\n new_opt.nc = nc\n torch.manualSeed(new_opt.manualSeed)\n local data_loader = paths.dofile('../data/data.lua')\n new_opt.phase = new_opt.phase .. name\n local data = data_loader.new(new_opt.nThreads, new_opt)\n print(\"Dataset Size \" .. name .. \": \", data:size())\n\n torch.setdefaulttensortype(tensortype)\n return data\nend\n\nreturn data_util\n" }, { "path": "data/dataset.lua", "content": "--[[\n Copyright (c) 2015-present, Facebook, Inc.\n All rights reserved.\n\n This source code is licensed under the BSD-style license found in the\n LICENSE file in the root directory of this source tree. An additional grant\n of patent rights can be found in the PATENTS file in the same directory.\n]]--\n\nrequire 'torch'\ntorch.setdefaulttensortype('torch.FloatTensor')\nlocal ffi = require 'ffi'\nlocal class = require('pl.class')\nlocal dir = require 'pl.dir'\nlocal tablex = require 'pl.tablex'\nlocal argcheck = require 'argcheck'\nrequire 'sys'\nrequire 'xlua'\nrequire 'image'\n\nlocal dataset = torch.class('dataLoader')\n\nlocal initcheck = argcheck{\n pack=true,\n help=[[\n A dataset class for images in a flat folder structure (folder-name is class-name).\n Optimized for extremely large datasets (upwards of 14 million images).\n Tested only on Linux (as it uses command-line linux utilities to scale up)\n]],\n {check=function(paths)\n local out = true;\n for k,v in ipairs(paths) do\n if type(v) ~= 'string' then\n print('paths can only be of string input');\n out = false\n end\n end\n return out\n end,\n name=\"paths\",\n type=\"table\",\n help=\"Multiple paths of directories with images\"},\n\n {name=\"sampleSize\",\n type=\"table\",\n help=\"a consistent sample size to resize the images\"},\n\n {name=\"split\",\n type=\"number\",\n help=\"Percentage of split to go to Training\"\n },\n {name=\"serial_batches\",\n type=\"number\",\n help=\"if randomly sample training images\"},\n\n {name=\"samplingMode\",\n type=\"string\",\n help=\"Sampling mode: random | balanced \",\n default = \"balanced\"},\n\n {name=\"verbose\",\n type=\"boolean\",\n help=\"Verbose mode during initialization\",\n default = false},\n\n {name=\"loadSize\",\n type=\"table\",\n help=\"a size to load the images to, initially\",\n opt = true},\n\n {name=\"forceClasses\",\n type=\"table\",\n help=\"If you want this loader to map certain classes to certain indices, \"\n .. \"pass a classes table that has {classname : classindex} pairs.\"\n .. \" For example: {3 : 'dog', 5 : 'cat'}\"\n .. \"This function is very useful when you want two loaders to have the same \"\n .. \"class indices (trainLoader/testLoader for example)\",\n opt = true},\n\n {name=\"sampleHookTrain\",\n type=\"function\",\n help=\"applied to sample during training(ex: for lighting jitter). \"\n .. \"It takes the image path as input\",\n opt = true},\n\n {name=\"sampleHookTest\",\n type=\"function\",\n help=\"applied to sample during testing\",\n opt = true},\n}\n\nfunction dataset:__init(...)\n\n -- argcheck\n local args = initcheck(...)\n print(args)\n for k,v in pairs(args) do self[k] = v end\n\n if not self.loadSize then self.loadSize = self.sampleSize; end\n\n if not self.sampleHookTrain then self.sampleHookTrain = self.defaultSampleHook end\n if not self.sampleHookTest then self.sampleHookTest = self.defaultSampleHook end\n self.image_count = 1\n-- print('image_count_init', self.image_count)\n -- find class names\n self.classes = {}\n local classPaths = {}\n if self.forceClasses then\n for k,v in pairs(self.forceClasses) do\n self.classes[k] = v\n classPaths[k] = {}\n end\n end\n local function tableFind(t, o) for k,v in pairs(t) do if v == o then return k end end end\n -- loop over each paths folder, get list of unique class names,\n -- also store the directory paths per class\n -- for each class,\n for k,path in ipairs(self.paths) do\n-- print('path', path)\n local dirs = {} -- hack\n dirs[1] = path\n-- local dirs = dir.getdirectories(path);\n for k,dirpath in ipairs(dirs) do\n local class = paths.basename(dirpath)\n local idx = tableFind(self.classes, class)\n-- print(class)\n-- print(idx)\n if not idx then\n table.insert(self.classes, class)\n idx = #self.classes\n classPaths[idx] = {}\n end\n if not tableFind(classPaths[idx], dirpath) then\n table.insert(classPaths[idx], dirpath);\n end\n end\n end\n\n self.classIndices = {}\n for k,v in ipairs(self.classes) do\n self.classIndices[v] = k\n end\n\n -- define command-line tools, try your best to maintain OSX compatibility\n local wc = 'wc'\n local cut = 'cut'\n local find = 'find -H' -- if folder name is symlink, do find inside it after dereferencing\n\n if ffi.os == 'OSX' then\n wc = 'gwc'\n cut = 'gcut'\n find = 'gfind'\n end\n ----------------------------------------------------------------------\n -- Options for the GNU find command\n local extensionList = {'jpg', 'png','JPG','PNG','JPEG', 'ppm', 'PPM', 'bmp', 'BMP'}\n local findOptions = ' -iname \"*.' .. extensionList[1] .. '\"'\n for i=2,#extensionList do\n findOptions = findOptions .. ' -o -iname \"*.' .. extensionList[i] .. '\"'\n end\n\n -- find the image path names\n self.imagePath = torch.CharTensor() -- path to each image in dataset\n self.imageClass = torch.LongTensor() -- class index of each image (class index in self.classes)\n self.classList = {} -- index of imageList to each image of a particular class\n self.classListSample = self.classList -- the main list used when sampling data\n\n print('running \"find\" on each class directory, and concatenate all'\n .. ' those filenames into a single file containing all image paths for a given class')\n -- so, generates one file per class\n local classFindFiles = {}\n for i=1,#self.classes do\n classFindFiles[i] = os.tmpname()\n end\n local combinedFindList = os.tmpname();\n\n local tmpfile = os.tmpname()\n local tmphandle = assert(io.open(tmpfile, 'w'))\n -- iterate over classes\n for i, class in ipairs(self.classes) do\n -- iterate over classPaths\n for j,path in ipairs(classPaths[i]) do\n local command = find .. ' \"' .. path .. '\" ' .. findOptions\n .. ' >>\"' .. classFindFiles[i] .. '\" \\n'\n tmphandle:write(command)\n end\n end\n io.close(tmphandle)\n os.execute('bash ' .. tmpfile)\n os.execute('rm -f ' .. tmpfile)\n\n print('now combine all the files to a single large file')\n local tmpfile = os.tmpname()\n local tmphandle = assert(io.open(tmpfile, 'w'))\n -- concat all finds to a single large file in the order of self.classes\n for i=1,#self.classes do\n local command = 'cat \"' .. classFindFiles[i] .. '\" >>' .. combinedFindList .. ' \\n'\n tmphandle:write(command)\n end\n io.close(tmphandle)\n os.execute('bash ' .. tmpfile)\n os.execute('rm -f ' .. tmpfile)\n\n --==========================================================================\n print('load the large concatenated list of sample paths to self.imagePath')\n local cmd = wc .. \" -L '\"\n .. combinedFindList .. \"' |\"\n .. cut .. \" -f1 -d' '\"\n print('cmd..' .. cmd)\n local maxPathLength = tonumber(sys.fexecute(wc .. \" -L '\"\n .. combinedFindList .. \"' |\"\n .. cut .. \" -f1 -d' '\")) + 1\n local length = tonumber(sys.fexecute(wc .. \" -l '\"\n .. combinedFindList .. \"' |\"\n .. cut .. \" -f1 -d' '\"))\n assert(length > 0, \"Could not find any image file in the given input paths\")\n assert(maxPathLength > 0, \"paths of files are length 0?\")\n self.imagePath:resize(length, maxPathLength):fill(0)\n local s_data = self.imagePath:data()\n local count = 0\n for line in io.lines(combinedFindList) do\n ffi.copy(s_data, line)\n s_data = s_data + maxPathLength\n if self.verbose and count % 10000 == 0 then\n xlua.progress(count, length)\n end;\n count = count + 1\n end\n\n self.numSamples = self.imagePath:size(1)\n if self.verbose then print(self.numSamples .. ' samples found.') end\n --==========================================================================\n print('Updating classList and imageClass appropriately')\n self.imageClass:resize(self.numSamples)\n local runningIndex = 0\n for i=1,#self.classes do\n if self.verbose then xlua.progress(i, #(self.classes)) end\n local length = tonumber(sys.fexecute(wc .. \" -l '\"\n .. classFindFiles[i] .. \"' |\"\n .. cut .. \" -f1 -d' '\"))\n if length == 0 then\n error('Class has zero samples')\n else\n self.classList[i] = torch.linspace(runningIndex + 1, runningIndex + length, length):long()\n self.imageClass[{{runningIndex + 1, runningIndex + length}}]:fill(i)\n end\n runningIndex = runningIndex + length\n end\n\n --==========================================================================\n -- clean up temporary files\n print('Cleaning up temporary files')\n local tmpfilelistall = ''\n for i=1,#(classFindFiles) do\n tmpfilelistall = tmpfilelistall .. ' \"' .. classFindFiles[i] .. '\"'\n if i % 1000 == 0 then\n os.execute('rm -f ' .. tmpfilelistall)\n tmpfilelistall = ''\n end\n end\n os.execute('rm -f ' .. tmpfilelistall)\n os.execute('rm -f \"' .. combinedFindList .. '\"')\n --==========================================================================\n\n if self.split == 100 then\n self.testIndicesSize = 0\n else\n print('Splitting training and test sets to a ratio of '\n .. self.split .. '/' .. (100-self.split))\n self.classListTrain = {}\n self.classListTest = {}\n self.classListSample = self.classListTrain\n local totalTestSamples = 0\n -- split the classList into classListTrain and classListTest\n for i=1,#self.classes do\n local list = self.classList[i]\n local count = self.classList[i]:size(1)\n local splitidx = math.floor((count * self.split / 100) + 0.5) -- +round\n local perm = torch.randperm(count)\n self.classListTrain[i] = torch.LongTensor(splitidx)\n for j=1,splitidx do\n self.classListTrain[i][j] = list[perm[j]]\n end\n if splitidx == count then -- all samples were allocated to train set\n self.classListTest[i] = torch.LongTensor()\n else\n self.classListTest[i] = torch.LongTensor(count-splitidx)\n totalTestSamples = totalTestSamples + self.classListTest[i]:size(1)\n local idx = 1\n for j=splitidx+1,count do\n self.classListTest[i][idx] = list[perm[j]]\n idx = idx + 1\n end\n end\n end\n -- Now combine classListTest into a single tensor\n self.testIndices = torch.LongTensor(totalTestSamples)\n self.testIndicesSize = totalTestSamples\n local tdata = self.testIndices:data()\n local tidx = 0\n for i=1,#self.classes do\n local list = self.classListTest[i]\n if list:dim() ~= 0 then\n local ldata = list:data()\n for j=0,list:size(1)-1 do\n tdata[tidx] = ldata[j]\n tidx = tidx + 1\n end\n end\n end\n end\nend\n\n-- size(), size(class)\nfunction dataset:size(class, list)\n list = list or self.classList\n if not class then\n return self.numSamples\n elseif type(class) == 'string' then\n return list[self.classIndices[class]]:size(1)\n elseif type(class) == 'number' then\n return list[class]:size(1)\n end\nend\n\n-- getByClass\nfunction dataset:getByClass(class)\n local index = 0\n if self.serial_batches == 1 then\n index = math.fmod(self.image_count-1, self.classListSample[class]:nElement())+1\n self.image_count = self.image_count +1\n else\n index = math.ceil(torch.uniform() * self.classListSample[class]:nElement())\n end\n\n local imgpath = ffi.string(torch.data(self.imagePath[self.classListSample[class][index]]))\n return self:sampleHookTrain(imgpath), imgpath\nend\n\n-- converts a table of samples (and corresponding labels) to a clean tensor\nlocal function tableToOutput(self, dataTable, scalarTable)\n local data, scalarLabels, labels\n if opt.resize_or_crop == 'crop' or opt.resize_or_crop == 'scale_width' or opt.resize_or_crop == 'scale_height' then\n assert(#scalarTable == 1)\n data = torch.Tensor(1,\n dataTable[1]:size(1), dataTable[1]:size(2), dataTable[1]:size(3))\n data[1]:copy(dataTable[1])\n scalarLabels = torch.LongTensor(#scalarTable):fill(-1111)\n else\n local quantity = #scalarTable\n data = torch.Tensor(quantity,\n self.sampleSize[1], self.sampleSize[2], self.sampleSize[3])\n scalarLabels = torch.LongTensor(quantity):fill(-1111)\n for i=1,#dataTable do\n data[i]:copy(dataTable[i])\n scalarLabels[i] = scalarTable[i]\n end\n end\n return data, scalarLabels\nend\n\n-- sampler, samples from the training set.\nfunction dataset:sample(quantity)\n assert(quantity)\n local dataTable = {}\n local scalarTable = {}\n local samplePaths = {}\n for i=1,quantity do\n local class = torch.random(1, #self.classes)\n local out, imgpath = self:getByClass(class)\n table.insert(dataTable, out)\n table.insert(scalarTable, class)\n samplePaths[i] = imgpath\n end\n\n local data, scalarLabels = tableToOutput(self, dataTable, scalarTable)\n return data, scalarLabels, samplePaths-- filePaths\nend\n\nfunction dataset:get(i1, i2)\n local indices = torch.range(i1, i2);\n local quantity = i2 - i1 + 1;\n assert(quantity > 0)\n -- now that indices has been initialized, get the samples\n local dataTable = {}\n local scalarTable = {}\n for i=1,quantity do\n -- load the sample\n local imgpath = ffi.string(torch.data(self.imagePath[indices[i]]))\n local out = self:sampleHookTest(imgpath)\n table.insert(dataTable, out)\n table.insert(scalarTable, self.imageClass[indices[i]])\n end\n local data, scalarLabels = tableToOutput(self, dataTable, scalarTable)\n return data, scalarLabels\nend\n\nreturn dataset\n" }, { "path": "data/donkey_folder.lua", "content": "\n--[[\n This data loader is a modified version of the one from dcgan.torch\n (see https://github.com/soumith/dcgan.torch/blob/master/data/donkey_folder.lua).\n Copyright (c) 2016, Deepak Pathak [See LICENSE file for details]\n Copyright (c) 2015-present, Facebook, Inc.\n All rights reserved.\n This source code is licensed under the BSD-style license found in the\n LICENSE file in the root directory of this source tree. An additional grant\n of patent rights can be found in the PATENTS file in the same directory.\n]]--\n\nrequire 'image'\npaths.dofile('dataset.lua')\n-- This file contains the data-loading logic and details.\n-- It is run by each data-loader thread.\n------------------------------------------\n-------- COMMON CACHES and PATHS\n-- Check for existence of opt.data\nif opt.DATA_ROOT then\n opt.data = paths.concat(opt.DATA_ROOT, opt.phase)\nelse\n print(os.getenv('DATA_ROOT'))\n opt.data = paths.concat(os.getenv('DATA_ROOT'), opt.phase)\nend\n\nif not paths.dirp(opt.data) then\n error('Did not find directory: ' .. opt.data)\nend\n\n-- a cache file of the training metadata (if doesnt exist, will be created)\nlocal cache_prefix = opt.data:gsub('/', '_')\nos.execute(('mkdir -p %s'):format(opt.cache_dir))\nlocal trainCache = paths.concat(opt.cache_dir, cache_prefix .. '_trainCache.t7')\n\n--------------------------------------------------------------------------------------------\nlocal input_nc = opt.nc -- input channels\nlocal loadSize = {input_nc, opt.loadSize}\nlocal sampleSize = {input_nc, opt.fineSize}\n\nlocal function loadImage(path)\n local input = image.load(path, 3, 'float')\n local h = input:size(2)\n local w = input:size(3)\n\n local imA = image.crop(input, 0, 0, w/2, h)\n imA = image.scale(imA, loadSize[2], loadSize[2])\n local imB = image.crop(input, w/2, 0, w, h)\n imB = image.scale(imB, loadSize[2], loadSize[2])\n\n local perm = torch.LongTensor{3, 2, 1}\n imA = imA:index(1, perm)\n imA = imA:mul(2):add(-1)\n imB = imB:index(1, perm)\n imB = imB:mul(2):add(-1)\n\n assert(imA:max()<=1,\"A: badly scaled inputs\")\n assert(imA:min()>=-1,\"A: badly scaled inputs\")\n assert(imB:max()<=1,\"B: badly scaled inputs\")\n assert(imB:min()>=-1,\"B: badly scaled inputs\")\n\n\n local oW = sampleSize[2]\n local oH = sampleSize[2]\n local iH = imA:size(2)\n local iW = imA:size(3)\n\n if iH~=oH then\n h1 = math.ceil(torch.uniform(1e-2, iH-oH))\n end\n\n if iW~=oW then\n w1 = math.ceil(torch.uniform(1e-2, iW-oW))\n end\n if iH ~= oH or iW ~= oW then\n imA = image.crop(imA, w1, h1, w1 + oW, h1 + oH)\n imB = image.crop(imB, w1, h1, w1 + oW, h1 + oH)\n end\n\n if opt.flip == 1 and torch.uniform() > 0.5 then\n imA = image.hflip(imA)\n imB = image.hflip(imB)\n end\n\n local concatenated = torch.cat(imA,imB,1)\n\n return concatenated\nend\n\n\nlocal function loadSingleImage(path)\n local im = image.load(path, input_nc, 'float')\n if opt.resize_or_crop == 'resize_and_crop' then\n im = image.scale(im, loadSize[2], loadSize[2])\n end\n if input_nc == 3 then\n local perm = torch.LongTensor{3, 2, 1}\n im = im:index(1, perm)--:mul(256.0): brg, rgb\n im = im:mul(2):add(-1)\n end\n assert(im:max()<=1,\"A: badly scaled inputs\")\n assert(im:min()>=-1,\"A: badly scaled inputs\")\n\n local oW = sampleSize[2]\n local oH = sampleSize[2]\n local iH = im:size(2)\n local iW = im:size(3)\n if (opt.resize_or_crop == 'resize_and_crop' ) then\n local h1, w1 = 0, 0\n if iH~=oH then\n h1 = math.ceil(torch.uniform(1e-2, iH-oH))\n end\n if iW~=oW then\n w1 = math.ceil(torch.uniform(1e-2, iW-oW))\n end\n if iH ~= oH or iW ~= oW then\n im = image.crop(im, w1, h1, w1 + oW, h1 + oH)\n end\n elseif (opt.resize_or_crop == 'combined') then\n local sH = math.min(math.ceil(oH * torch.uniform(1+1e-2, 2.0-1e-2)), iH-1e-2)\n local sW = math.min(math.ceil(oW * torch.uniform(1+1e-2, 2.0-1e-2)), iW-1e-2)\n local h1 = math.ceil(torch.uniform(1e-2, iH-sH))\n local w1 = math.ceil(torch.uniform(1e-2, iW-sW))\n im = image.crop(im, w1, h1, w1 + sW, h1 + sH)\n im = image.scale(im, oW, oH)\n elseif (opt.resize_or_crop == 'crop') then\n local w = math.min(math.min(oH, iH),iW)\n w = math.floor(w/4)*4\n local x = math.floor(torch.uniform(0, iW - w))\n local y = math.floor(torch.uniform(0, iH - w))\n im = image.crop(im, x, y, x+w, y+w)\n elseif (opt.resize_or_crop == 'scale_width') then\n w = oW\n h = torch.floor(iH * oW/iW)\n im = image.scale(im, w, h)\n elseif (opt.resize_or_crop == 'scale_height') then\n h = oH\n w = torch.floor(iW * oH / iH)\n im = image.scale(im, w, h)\n end\n\n if opt.flip == 1 and torch.uniform() > 0.5 then\n im = image.hflip(im)\n end\n\n return im\n\nend\n\n-- channel-wise mean and std. Calculate or load them from disk later in the script.\nlocal mean,std\n--------------------------------------------------------------------------------\n-- Hooks that are used for each image that is loaded\n\n-- function to load the image, jitter it appropriately (random crops etc.)\nlocal trainHook_singleimage = function(self, path)\n collectgarbage()\n -- print('load single image')\n local im = loadSingleImage(path)\n return im\nend\n\n-- function that loads images that have juxtaposition\n-- of two images from two domains\nlocal trainHook_doubleimage = function(self, path)\n -- print('load double image')\n collectgarbage()\n\n local im = loadImage(path)\n return im\nend\n\n\nif opt.align_data > 0 then\n sample_nc = input_nc*2\n trainHook = trainHook_doubleimage\nelse\n sample_nc = input_nc\n trainHook = trainHook_singleimage\nend\n\ntrainLoader = dataLoader{\n paths = {opt.data},\n loadSize = {input_nc, loadSize[2], loadSize[2]},\n sampleSize = {sample_nc, sampleSize[2], sampleSize[2]},\n split = 100,\n serial_batches = opt.serial_batches,\n verbose = true\n }\n\ntrainLoader.sampleHookTrain = trainHook\ncollectgarbage()\n\n-- do some sanity checks on trainLoader\ndo\n local class = trainLoader.imageClass\n local nClasses = #trainLoader.classes\n assert(class:max() <= nClasses, \"class logic has error\")\n assert(class:min() >= 1, \"class logic has error\")\nend\n" }, { "path": "data/unaligned_data_loader.lua", "content": "--------------------------------------------------------------------------------\n-- Subclass of BaseDataLoader that provides data from two datasets.\n-- The samples from the datasets are not aligned.\n-- The datasets can have different sizes\n--------------------------------------------------------------------------------\nrequire 'data.base_data_loader'\n\nlocal class = require 'class'\ndata_util = paths.dofile('data_util.lua')\n\nUnalignedDataLoader = class('UnalignedDataLoader', 'BaseDataLoader')\n\nfunction UnalignedDataLoader:__init(conf)\n BaseDataLoader.__init(self, conf)\n conf = conf or {}\nend\n\nfunction UnalignedDataLoader:name()\n return 'UnalignedDataLoader'\nend\n\nfunction UnalignedDataLoader:Initialize(opt)\n opt.align_data = 0\n self.dataA = data_util.load_dataset('A', opt, opt.input_nc)\n self.dataB = data_util.load_dataset('B', opt, opt.output_nc)\nend\n\n-- actually fetches the data\n-- |return|: a table of two tables, each corresponding to\n-- the batch for dataset A and dataset B\nfunction UnalignedDataLoader:LoadBatchForAllDatasets()\n local batchA, pathA = self.dataA:getBatch()\n local batchB, pathB = self.dataB:getBatch()\n return batchA, batchB, pathA, pathB\nend\n\n-- returns the size of each dataset\nfunction UnalignedDataLoader:size(dataset)\n if dataset == 'A' then\n return self.dataA:size()\n end\n\n if dataset == 'B' then\n return self.dataB:size()\n end\n\n return math.max(self.dataA:size(), self.dataB:size())\n -- return the size of the largest dataset by default\nend\n" }, { "path": "examples/test_vangogh_style_on_ae_photos.sh", "content": "#!/bin/sh\n\n## This script download the dataset and pre-trained network,\n## and generates style transferred images.\n\n# Download the dataset. The downloaded dataset is stored in ./datasets/${DATASET_NAME}\nDATASET_NAME='ae_photos'\nbash ./datasets/download_dataset.sh $DATASET_NAME\n\n# Download the pre-trained model. The downloaded model is stored in ./models/${MODEL_NAME}_pretrained/latest_net_G.t7\nMODEL_NAME='style_vangogh'\nbash ./pretrained_models/download_model.sh $MODEL_NAME\n\n# Run style transfer using the downloaded dataset and model\nDATA_ROOT=./datasets/$DATASET_NAME name=${MODEL_NAME}_pretrained model=one_direction_test phase=test how_many='all' loadSize=256 fineSize=256 resize_or_crop='scale_width' th test.lua\n\nif [ $? == 0 ]; then\n echo \"The result can be viewed at ./results/${MODEL_NAME}_pretrained/latest_test/index.html\"\nfi\n" }, { "path": "examples/train_maps.sh", "content": "DB_NAME='maps'\nGPU_ID=1\nDISPLAY_ID=1\nNET_G=resnet_6blocks\nNET_D=basic\nMODEL=cycle_gan\nSAVE_EPOCH=5\nALIGN_DATA=0\nLAMBDA=10\nNF=64\n\n\nEXPR_NAME=${DB_NAME}_${MODEL}_${LAMBDA}\n\nCHECKPOINT_DIR=./checkpoints/\nLOG_FILE=${CHECKPOINT_DIR}${EXPR_NAME}/log.txt\nmkdir -p ${CHECKPOINT_DIR}${EXPR_NAME}\n\nDATA_ROOT=./datasets/$DB_NAME align_data=$ALIGN_DATA use_lsgan=1 \\\nwhich_direction='AtoB' display_plot=$PLOT pool_size=50 niter=100 niter_decay=100 \\\nwhich_model_netG=$NET_G which_model_netD=$NET_D model=$MODEL lr=0.0002 print_freq=200 lambda_A=$LAMBDA lambda_B=$LAMBDA \\\nloadSize=143 fineSize=128 gpu=$GPU_ID display_winsize=128 \\\nname=$EXPR_NAME flip=1 save_epoch_freq=$SAVE_EPOCH \\\ncontinue_train=0 display_id=$DISPLAY_ID \\\ncheckpoints_dir=$CHECKPOINT_DIR\\\n th train.lua | tee -a $LOG_FILE\n" }, { "path": "models/architectures.lua", "content": "require 'nngraph'\n\n\n----------------------------------------------------------------------------\nlocal function weights_init(m)\n local name = torch.type(m)\n if name:find('Convolution') then\n m.weight:normal(0.0, 0.02)\n m.bias:fill(0)\n elseif name:find('Normalization') then\n if m.weight then m.weight:normal(1.0, 0.02) end\n if m.bias then m.bias:fill(0) end\n end\nend\n\n\nnormalization = nil\n\nfunction set_normalization(norm)\nif norm == 'instance' then\n require 'util.InstanceNormalization'\n print('use InstanceNormalization')\n normalization = nn.InstanceNormalization\nelseif norm == 'batch' then\n print('use SpatialBatchNormalization')\n normalization = nn.SpatialBatchNormalization\nend\nend\n\nfunction defineG(input_nc, output_nc, ngf, which_model_netG, nz, arch)\n local netG = nil\n if which_model_netG == \"encoder_decoder\" then netG = defineG_encoder_decoder(input_nc, output_nc, ngf)\n elseif which_model_netG == \"unet128\" then netG = defineG_unet128(input_nc, output_nc, ngf)\n elseif which_model_netG == \"unet256\" then netG = defineG_unet256(input_nc, output_nc, ngf)\n elseif which_model_netG == \"resnet_6blocks\" then netG = defineG_resnet_6blocks(input_nc, output_nc, ngf)\n elseif which_model_netG == \"resnet_9blocks\" then netG = defineG_resnet_9blocks(input_nc, output_nc, ngf)\n else error(\"unsupported netG model\")\n end\n netG:apply(weights_init)\n\n return netG\nend\n\nfunction defineD(input_nc, ndf, which_model_netD, n_layers_D, use_sigmoid)\n local netD = nil\n if which_model_netD == \"basic\" then netD = defineD_basic(input_nc, ndf, use_sigmoid)\n elseif which_model_netD == \"imageGAN\" then netD = defineD_imageGAN(input_nc, ndf, use_sigmoid)\n elseif which_model_netD == \"n_layers\" then netD = defineD_n_layers(input_nc, ndf, n_layers_D, use_sigmoid)\n else error(\"unsupported netD model\")\n end\n netD:apply(weights_init)\n\n return netD\nend\n\nfunction defineG_encoder_decoder(input_nc, output_nc, ngf)\n -- input is (nc) x 256 x 256\n local e1 = - nn.SpatialConvolution(input_nc, ngf, 4, 4, 2, 2, 1, 1)\n -- input is (ngf) x 128 x 128\n local e2 = e1 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf, ngf * 2, 4, 4, 2, 2, 1, 1) - normalization(ngf * 2)\n -- input is (ngf * 2) x 64 x 64\n local e3 = e2 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 2, ngf * 4, 4, 4, 2, 2, 1, 1) - normalization(ngf * 4)\n -- input is (ngf * 4) x 32 x 32\n local e4 = e3 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 4, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 16 x 16\n local e5 = e4 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 8 x 8\n local e6 = e5 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 4 x 4\n local e7 = e6 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 2 x 2\n local e8 = e7 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) -- normalization(ngf * 8)\n -- input is (ngf * 8) x 1 x 1\n\n local d1 = e8 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8) - nn.Dropout(0.5)\n -- input is (ngf * 8) x 2 x 2\n local d2 = d1 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8) - nn.Dropout(0.5)\n -- input is (ngf * 8) x 4 x 4\n local d3 = d2 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8) - nn.Dropout(0.5)\n -- input is (ngf * 8) x 8 x 8\n local d4 = d3 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 16 x 16\n local d5 = d4 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8, ngf * 4, 4, 4, 2, 2, 1, 1) - normalization(ngf * 4)\n -- input is (ngf * 4) x 32 x 32\n local d6 = d5 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 4, ngf * 2, 4, 4, 2, 2, 1, 1) - normalization(ngf * 2)\n -- input is (ngf * 2) x 64 x 64\n local d7 = d6 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 2, ngf, 4, 4, 2, 2, 1, 1) - normalization(ngf)\n -- input is (ngf) x128 x 128\n local d8 = d7 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf, output_nc, 4, 4, 2, 2, 1, 1)\n -- input is (nc) x 256 x 256\n local o1 = d8 - nn.Tanh()\n\n local netG = nn.gModule({e1},{o1})\n return netG\nend\n\n\nfunction defineG_unet128(input_nc, output_nc, ngf)\n local netG = nil\n -- input is (nc) x 128 x 128\n local e1 = - nn.SpatialConvolution(input_nc, ngf, 4, 4, 2, 2, 1, 1)\n -- input is (ngf) x 64 x 64\n local e2 = e1 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf, ngf * 2, 4, 4, 2, 2, 1, 1) - normalization(ngf * 2)\n -- input is (ngf * 2) x 32 x 32\n local e3 = e2 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 2, ngf * 4, 4, 4, 2, 2, 1, 1) - normalization(ngf * 4)\n -- input is (ngf * 4) x 16 x 16\n local e4 = e3 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 4, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 8 x 8\n local e5 = e4 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 4 x 4\n local e6 = e5 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 2 x 2\n local e7 = e6 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) -- normalization(ngf * 8)\n -- input is (ngf * 8) x 1 x 1\n\n local d1_ = e7 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8) - nn.Dropout(0.5)\n -- input is (ngf * 8) x 2 x 2\n local d1 = {d1_,e6} - nn.JoinTable(2)\n local d2_ = d1 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8 * 2, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8) - nn.Dropout(0.5)\n -- input is (ngf * 8) x 4 x 4\n local d2 = {d2_,e5} - nn.JoinTable(2)\n local d3_ = d2 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8 * 2, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8) - nn.Dropout(0.5)\n -- input is (ngf * 8) x 8 x 8\n local d3 = {d3_,e4} - nn.JoinTable(2)\n local d4_ = d3 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8 * 2, ngf * 4, 4, 4, 2, 2, 1, 1) - normalization(ngf * 4)\n -- input is (ngf * 8) x 16 x 16\n local d4 = {d4_,e3} - nn.JoinTable(2)\n local d5_ = d4 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 4 * 2, ngf * 2, 4, 4, 2, 2, 1, 1) - normalization(ngf * 2)\n -- input is (ngf * 4) x 32 x 32\n local d5 = {d5_,e2} - nn.JoinTable(2)\n local d6_ = d5 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 2 * 2, ngf, 4, 4, 2, 2, 1, 1) - normalization(ngf)\n -- input is (ngf * 2) x 64 x 64\n local d6 = {d6_,e1} - nn.JoinTable(2)\n\n local d7 = d6 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 2, output_nc, 4, 4, 2, 2, 1, 1)\n -- input is (nc) x 128 x 128\n\n local o1 = d7 - nn.Tanh()\n local netG = nn.gModule({e1},{o1})\n return netG\nend\n\n\nfunction defineG_unet256(input_nc, output_nc, ngf)\n local netG = nil\n -- input is (nc) x 256 x 256\n local e1 = - nn.SpatialConvolution(input_nc, ngf, 4, 4, 2, 2, 1, 1)\n -- input is (ngf) x 128 x 128\n local e2 = e1 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf, ngf * 2, 4, 4, 2, 2, 1, 1) - normalization(ngf * 2)\n -- input is (ngf * 2) x 64 x 64\n local e3 = e2 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 2, ngf * 4, 4, 4, 2, 2, 1, 1) - normalization(ngf * 4)\n -- input is (ngf * 4) x 32 x 32\n local e4 = e3 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 4, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 16 x 16\n local e5 = e4 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 8 x 8\n local e6 = e5 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 4 x 4\n local e7 = e6 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 2 x 2\n local e8 = e7 - nn.LeakyReLU(0.2, true) - nn.SpatialConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) -- - normalization(ngf * 8)\n -- input is (ngf * 8) x 1 x 1\n\n local d1_ = e8 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8) - nn.Dropout(0.5)\n -- input is (ngf * 8) x 2 x 2\n local d1 = {d1_,e7} - nn.JoinTable(2)\n local d2_ = d1 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8 * 2, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8) - nn.Dropout(0.5)\n -- input is (ngf * 8) x 4 x 4\n local d2 = {d2_,e6} - nn.JoinTable(2)\n local d3_ = d2 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8 * 2, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8) - nn.Dropout(0.5)\n -- input is (ngf * 8) x 8 x 8\n local d3 = {d3_,e5} - nn.JoinTable(2)\n local d4_ = d3 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8 * 2, ngf * 8, 4, 4, 2, 2, 1, 1) - normalization(ngf * 8)\n -- input is (ngf * 8) x 16 x 16\n local d4 = {d4_,e4} - nn.JoinTable(2)\n local d5_ = d4 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 8 * 2, ngf * 4, 4, 4, 2, 2, 1, 1) - normalization(ngf * 4)\n -- input is (ngf * 4) x 32 x 32\n local d5 = {d5_,e3} - nn.JoinTable(2)\n local d6_ = d5 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 4 * 2, ngf * 2, 4, 4, 2, 2, 1, 1) - normalization(ngf * 2)\n -- input is (ngf * 2) x 64 x 64\n local d6 = {d6_,e2} - nn.JoinTable(2)\n local d7_ = d6 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 2 * 2, ngf, 4, 4, 2, 2, 1, 1) - normalization(ngf)\n -- input is (ngf) x128 x 128\n local d7 = {d7_,e1} - nn.JoinTable(2)\n local d8 = d7 - nn.ReLU(true) - nn.SpatialFullConvolution(ngf * 2, output_nc, 4, 4, 2, 2, 1, 1)\n -- input is (nc) x 256 x 256\n\n local o1 = d8 - nn.Tanh()\n local netG = nn.gModule({e1},{o1})\n return netG\nend\n\n--------------------------------------------------------------------------------\n-- Justin Johnson's model from https://github.com/jcjohnson/fast-neural-style/\n--------------------------------------------------------------------------------\n\nlocal function build_conv_block(dim, padding_type)\n local conv_block = nn.Sequential()\n local p = 0\n if padding_type == 'reflect' then\n conv_block:add(nn.SpatialReflectionPadding(1, 1, 1, 1))\n elseif padding_type == 'replicate' then\n conv_block:add(nn.SpatialReplicationPadding(1, 1, 1, 1))\n elseif padding_type == 'zero' then\n p = 1\n end\n conv_block:add(nn.SpatialConvolution(dim, dim, 3, 3, 1, 1, p, p))\n conv_block:add(normalization(dim))\n conv_block:add(nn.ReLU(true))\n if padding_type == 'reflect' then\n conv_block:add(nn.SpatialReflectionPadding(1, 1, 1, 1))\n elseif padding_type == 'replicate' then\n conv_block:add(nn.SpatialReplicationPadding(1, 1, 1, 1))\n end\n conv_block:add(nn.SpatialConvolution(dim, dim, 3, 3, 1, 1, p, p))\n conv_block:add(normalization(dim))\n return conv_block\nend\n\n\nlocal function build_res_block(dim, padding_type)\n local conv_block = build_conv_block(dim, padding_type)\n local res_block = nn.Sequential()\n local concat = nn.ConcatTable()\n concat:add(conv_block)\n concat:add(nn.Identity())\n \n res_block:add(concat):add(nn.CAddTable())\n return res_block\nend\n\nfunction defineG_resnet_6blocks(input_nc, output_nc, ngf)\n padding_type = 'reflect'\n local ks = 3\n local netG = nil\n local f = 7\n local p = (f - 1) / 2\n local data = -nn.Identity()\n local e1 = data - nn.SpatialReflectionPadding(p, p, p, p) - nn.SpatialConvolution(input_nc, ngf, f, f, 1, 1) - normalization(ngf) - nn.ReLU(true)\n local e2 = e1 - nn.SpatialConvolution(ngf, ngf*2, ks, ks, 2, 2, 1, 1) - normalization(ngf*2) - nn.ReLU(true)\n local e3 = e2 - nn.SpatialConvolution(ngf*2, ngf*4, ks, ks, 2, 2, 1, 1) - normalization(ngf*4) - nn.ReLU(true)\n local d1 = e3 - build_res_block(ngf*4, padding_type) - build_res_block(ngf*4, padding_type) - build_res_block(ngf*4, padding_type)\n - build_res_block(ngf*4, padding_type) - build_res_block(ngf*4, padding_type) - build_res_block(ngf*4, padding_type)\n local d2 = d1 - nn.SpatialFullConvolution(ngf*4, ngf*2, ks, ks, 2, 2, 1, 1,1,1) - normalization(ngf*2) - nn.ReLU(true)\n local d3 = d2 - nn.SpatialFullConvolution(ngf*2, ngf, ks, ks, 2, 2, 1, 1,1,1) - normalization(ngf) - nn.ReLU(true)\n local d4 = d3 - nn.SpatialReflectionPadding(p, p, p, p) - nn.SpatialConvolution(ngf, output_nc, f, f, 1, 1) - nn.Tanh()\n netG = nn.gModule({data},{d4})\n return netG\nend\n\nfunction defineG_resnet_9blocks(input_nc, output_nc, ngf)\n padding_type = 'reflect'\n local ks = 3\n local netG = nil\n local f = 7\n local p = (f - 1) / 2\n local data = -nn.Identity()\n local e1 = data - nn.SpatialReflectionPadding(p, p, p, p) - nn.SpatialConvolution(input_nc, ngf, f, f, 1, 1) - normalization(ngf) - nn.ReLU(true)\n local e2 = e1 - nn.SpatialConvolution(ngf, ngf*2, ks, ks, 2, 2, 1, 1) - normalization(ngf*2) - nn.ReLU(true)\n local e3 = e2 - nn.SpatialConvolution(ngf*2, ngf*4, ks, ks, 2, 2, 1, 1) - normalization(ngf*4) - nn.ReLU(true)\n local d1 = e3 - build_res_block(ngf*4, padding_type) - build_res_block(ngf*4, padding_type) - build_res_block(ngf*4, padding_type)\n - build_res_block(ngf*4, padding_type) - build_res_block(ngf*4, padding_type) - build_res_block(ngf*4, padding_type)\n - build_res_block(ngf*4, padding_type) - build_res_block(ngf*4, padding_type) - build_res_block(ngf*4, padding_type)\n local d2 = d1 - nn.SpatialFullConvolution(ngf*4, ngf*2, ks, ks, 2, 2, 1, 1,1,1) - normalization(ngf*2) - nn.ReLU(true)\n local d3 = d2 - nn.SpatialFullConvolution(ngf*2, ngf, ks, ks, 2, 2, 1, 1,1,1) - normalization(ngf) - nn.ReLU(true)\n local d4 = d3 - nn.SpatialReflectionPadding(p, p, p, p) - nn.SpatialConvolution(ngf, output_nc, f, f, 1, 1) - nn.Tanh()\n netG = nn.gModule({data},{d4})\n return netG\nend\n\nfunction defineD_imageGAN(input_nc, ndf, use_sigmoid)\n local netD = nn.Sequential()\n\n -- input is (nc) x 256 x 256\n netD:add(nn.SpatialConvolution(input_nc, ndf, 4, 4, 2, 2, 1, 1))\n netD:add(nn.LeakyReLU(0.2, true))\n -- state size: (ndf) x 128 x 128\n netD:add(nn.SpatialConvolution(ndf, ndf * 2, 4, 4, 2, 2, 1, 1))\n netD:add(nn.SpatialBatchNormalization(ndf * 2)):add(nn.LeakyReLU(0.2, true))\n -- state size: (ndf*2) x 64 x 64\n netD:add(nn.SpatialConvolution(ndf * 2, ndf*4, 4, 4, 2, 2, 1, 1))\n netD:add(nn.SpatialBatchNormalization(ndf * 4)):add(nn.LeakyReLU(0.2, true))\n -- state size: (ndf*4) x 32 x 32\n netD:add(nn.SpatialConvolution(ndf * 4, ndf * 8, 4, 4, 2, 2, 1, 1))\n netD:add(nn.SpatialBatchNormalization(ndf * 8)):add(nn.LeakyReLU(0.2, true))\n -- state size: (ndf*8) x 16 x 16\n netD:add(nn.SpatialConvolution(ndf * 8, ndf * 8, 4, 4, 2, 2, 1, 1))\n netD:add(nn.SpatialBatchNormalization(ndf * 8)):add(nn.LeakyReLU(0.2, true))\n -- state size: (ndf*8) x 8 x 8\n netD:add(nn.SpatialConvolution(ndf * 8, ndf * 8, 4, 4, 2, 2, 1, 1))\n netD:add(nn.SpatialBatchNormalization(ndf * 8)):add(nn.LeakyReLU(0.2, true))\n -- state size: (ndf*8) x 4 x 4\n netD:add(nn.SpatialConvolution(ndf * 8, 1, 4, 4, 2, 2, 1, 1))\n -- state size: 1 x 1 x 1\n if use_sigmoid then\n netD:add(nn.Sigmoid())\n end\n\n\treturn netD\nend\n\n\n\nfunction defineD_basic(input_nc, ndf, use_sigmoid)\n n_layers = 3\n return defineD_n_layers(input_nc, ndf, n_layers, use_sigmoid)\nend\n\n-- rf=1\nfunction defineD_pixelGAN(input_nc, ndf, use_sigmoid)\n\n local netD = nn.Sequential()\n\n -- input is (nc) x 256 x 256\n netD:add(nn.SpatialConvolution(input_nc, ndf, 1, 1, 1, 1, 0, 0))\n netD:add(nn.LeakyReLU(0.2, true))\n -- state size: (ndf) x 256 x 256\n netD:add(nn.SpatialConvolution(ndf, ndf * 2, 1, 1, 1, 1, 0, 0))\n netD:add(normalization(ndf * 2)):add(nn.LeakyReLU(0.2, true))\n -- state size: (ndf*2) x 256 x 256\n netD:add(nn.SpatialConvolution(ndf * 2, 1, 1, 1, 1, 1, 0, 0))\n -- state size: 1 x 256 x 256\n if use_sigmoid then\n netD:add(nn.Sigmoid())\n -- state size: 1 x 30 x 30\n end\n\n return netD\nend\n\n-- if n=0, then use pixelGAN (rf=1)\n-- else rf is 16 if n=1\n-- 34 if n=2\n-- 70 if n=3\n-- 142 if n=4\n-- 286 if n=5\n-- 574 if n=6\nfunction defineD_n_layers(input_nc, ndf, n_layers, use_sigmoid, kw, dropout_ratio)\n\n if dropout_ratio == nil then\n dropout_ratio = 0.0\n end\n\n if kw == nil then\n\tkw = 4\n end\n padw = math.ceil((kw-1)/2)\n\n if n_layers==0 then\n return defineD_pixelGAN(input_nc, ndf, use_sigmoid)\n else\n\n local netD = nn.Sequential()\n\n -- input is (nc) x 256 x 256\n -- print('input_nc', input_nc)\n netD:add(nn.SpatialConvolution(input_nc, ndf, kw, kw, 2, 2, padw, padw))\n netD:add(nn.LeakyReLU(0.2, true))\n\n local nf_mult = 1\n local nf_mult_prev = 1\n for n = 1, n_layers-1 do\n nf_mult_prev = nf_mult\n nf_mult = math.min(2^n,8)\n netD:add(nn.SpatialConvolution(ndf * nf_mult_prev, ndf * nf_mult, kw, kw, 2, 2, padw,padw))\n netD:add(normalization(ndf * nf_mult)):add(nn.Dropout(dropout_ratio))\n netD:add(nn.LeakyReLU(0.2, true))\n end\n\n -- state size: (ndf*M) x N x N\n nf_mult_prev = nf_mult\n nf_mult = math.min(2^n_layers,8)\n netD:add(nn.SpatialConvolution(ndf * nf_mult_prev, ndf * nf_mult, kw, kw, 1, 1, padw, padw))\n netD:add(normalization(ndf * nf_mult)):add(nn.LeakyReLU(0.2, true))\n -- state size: (ndf*M*2) x (N-1) x (N-1)\n netD:add(nn.SpatialConvolution(ndf * nf_mult, 1, kw, kw, 1, 1, padw,padw))\n -- state size: 1 x (N-2) x (N-2)\n if use_sigmoid then\n netD:add(nn.Sigmoid())\n end\n -- state size: 1 x (N-2) x (N-2)\n return netD\n end\nend\n" }, { "path": "models/base_model.lua", "content": "--------------------------------------------------------------------------------\n-- Base Class for Providing Models\n--------------------------------------------------------------------------------\n\nlocal class = require 'class'\n\nBaseModel = class('BaseModel')\n\nfunction BaseModel:__init(conf)\n conf = conf or {}\nend\n\n-- Returns the name of the model\nfunction BaseModel:model_name()\n return 'DoesNothingModel'\nend\n\n-- Defines models and networks\nfunction BaseModel:Initialize(opt)\n models = {}\n return models\nend\n\n-- Runs the forward pass of the network\nfunction BaseModel:Forward(input, opt)\n output = {}\n return output\nend\n\n-- Runs the backprop gradient descent\n-- Corresponds to a single batch of data\nfunction BaseModel:OptimizeParameters(opt)\nend\n\n-- This function can be used to reset momentum after each epoch\nfunction BaseModel:RefreshParameters(opt)\nend\n\n-- This function can be used to reset momentum after each epoch\nfunction BaseModel:UpdateLearningRate(opt)\nend\n-- Save the current model to the file system\nfunction BaseModel:Save(prefix, opt)\nend\n\n-- returns a string that describes the current errors\nfunction BaseModel:GetCurrentErrorDescription()\n return \"No Error exists in BaseModel\"\nend\n\n-- returns current errors\nfunction BaseModel:GetCurrentErrors(opt)\n return {}\nend\n\n-- returns a table of image/label pairs that describe\n-- the current results.\n-- |return|: a table of table. List of image/label pairs\nfunction BaseModel:GetCurrentVisuals(opt, size)\n return {}\nend\n\n-- returns a string that describes the display plot configuration \nfunction BaseModel:DisplayPlot(opt)\n return {}\nend\n" }, { "path": "models/bigan_model.lua", "content": "local class = require 'class'\nrequire 'models.base_model'\nrequire 'models.architectures'\nrequire 'util.image_pool'\nutil = paths.dofile('../util/util.lua')\ncontent = paths.dofile('../util/content_loss.lua')\n\nBiGANModel = class('BiGANModel', 'BaseModel')\n\nfunction BiGANModel:__init(conf)\n BaseModel.__init(self, conf)\n conf = conf or {}\nend\n\nfunction BiGANModel:model_name()\n return 'BiGANModel'\nend\n\nfunction BiGANModel:InitializeStates(use_wgan)\n optimState = {learningRate=opt.lr, beta1=opt.beta1,}\n return optimState\nend\n-- Defines models and networks\nfunction BiGANModel:Initialize(opt)\n if opt.test == 0 then\n self.realABPool = ImagePool(opt.pool_size)\n self.fakeABPool = ImagePool(opt.pool_size)\n end\n -- define tensors\n local d_input_nc = opt.input_nc + opt.output_nc\n self.real_AB = torch.Tensor(opt.batchSize, d_input_nc, opt.fineSize, opt.fineSize)\n self.fake_AB = torch.Tensor(opt.batchSize, d_input_nc, opt.fineSize, opt.fineSize)\n -- load/define models\n self.criterionGAN = nn.MSECriterion()\n\n local netG, netE, netD = nil, nil, nil\n if opt.continue_train == 1 then\n if opt.test == 1 then -- which_epoch option exists in test mode\n netG = util.load_test_model('G', opt)\n netE = util.load_test_model('E', opt)\n netD = util.load_test_model('D', opt)\n else\n netG = util.load_model('G', opt)\n netE = util.load_model('E', opt)\n netD = util.load_model('D', opt)\n end\n else\n -- netG_test = defineG(opt.input_nc, opt.output_nc, opt.ngf, \"resnet_unet\", opt.arch)\n -- os.exit()\n netD = defineD(d_input_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, false) -- no sigmoid layer\n print('netD...', netD)\n netG = defineG(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.arch)\n print('netG...', netG)\n netE = defineG(opt.output_nc, opt.input_nc, opt.ngf, opt.which_model_netG, opt.arch)\n print('netE...', netE)\n\n end\n\n self.netD = netD\n self.netG = netG\n self.netE = netE\n\n -- define real/fake labels\n netD_output_size = self.netD:forward(self.real_AB):size()\n self.fake_label = torch.Tensor(netD_output_size):fill(0.0)\n self.real_label = torch.Tensor(netD_output_size):fill(1.0) -- no soft smoothing\n\n self.optimStateD = self:InitializeStates()\n self.optimStateG = self:InitializeStates()\n self.optimStateE = self:InitializeStates()\n self.A_idx = {{}, {1, opt.input_nc}, {}, {}}\n self.B_idx = {{}, {opt.input_nc+1, opt.input_nc+opt.output_nc}, {}, {}}\n self:RefreshParameters()\n\n print('---------- # Learnable Parameters --------------')\n print(('G = %d'):format(self.parametersG:size(1)))\n print(('E = %d'):format(self.parametersE:size(1)))\n print(('D = %d'):format(self.parametersD:size(1)))\n print('------------------------------------------------')\n -- os.exit()\nend\n\n-- Runs the forward pass of the network and\n-- saves the result to member variables of the class\nfunction BiGANModel:Forward(input, opt)\n if opt.which_direction == 'BtoA' then\n \tlocal temp = input.real_A\n \tinput.real_A = input.real_B\n \tinput.real_B = temp\n end\n self.real_AB[self.A_idx]:copy(input.real_A)\n self.fake_AB[self.B_idx]:copy(input.real_B)\n self.real_A = self.real_AB[self.A_idx]\n self.real_B = self.fake_AB[self.B_idx]\n self.fake_B = self.netG:forward(self.real_A):clone()\n self.fake_A = self.netE:forward(self.real_B):clone()\n self.real_AB[self.B_idx]:copy(self.fake_B) -- real_AB: real_A, fake_B -> real_label\n self.fake_AB[self.A_idx]:copy(self.fake_A) -- fake_AB: fake_A, real_B -> fake_label\n -- if opt.test == 0 then\n -- self.real_AB = self.realABPool:Query(self.real_AB) -- batch history\n -- self.fake_AB = self.fakeABPool:Query(self.fake_AB) -- batch history\n -- end\nend\n\n-- create closure to evaluate f(X) and df/dX of discriminator\nfunction BiGANModel:fDx_basic(x, gradParams, netD, real_AB, fake_AB, opt)\n util.BiasZero(netD)\n gradParams:zero()\n -- Real log(D_A(B))\n local output = netD:forward(real_AB):clone()\n local errD_real = self.criterionGAN:forward(output, self.real_label)\n local df_do = self.criterionGAN:backward(output, self.real_label)\n netD:backward(real_AB, df_do)\n -- Fake + log(1 - D_A(G(A)))\n output = netD:forward(fake_AB):clone()\n local errD_fake = self.criterionGAN:forward(output, self.fake_label)\n local df_do2 = self.criterionGAN:backward(output, self.fake_label)\n netD:backward(fake_AB, df_do2)\n -- Compute loss\n local errD = (errD_real + errD_fake) / 2.0\n return errD, gradParams\nend\n\n\nfunction BiGANModel:fDx(x, opt)\n -- use image pool that stores the old fake images\n real_AB = self.realABPool:Query(self.real_AB)\n fake_AB = self.fakeABPool:Query(self.fake_AB)\n self.errD, gradParams = self:fDx_basic(x, self.gradParametersD, self.netD, real_AB, fake_AB, opt)\n return self.errD, gradParams\nend\n\n\n\nfunction BiGANModel:fGx_basic(x, netG, netD, gradParametersG, opt)\n util.BiasZero(netG)\n util.BiasZero(netD)\n gradParametersG:zero()\n\n -- First. G(A) should fake the discriminator\n local output = netD:forward(self.real_AB):clone()\n local errG = self.criterionGAN:forward(output, self.fake_label)\n local dgan_loss_dd = self.criterionGAN:backward(output, self.fake_label)\n local dgan_loss_do = netD:updateGradInput(self.real_AB, dgan_loss_dd)\n netG:backward(self.real_A, dgan_loss_do[self.B_idx]) -- real_AB: real_A, fake_B -> real_label\n return gradParametersG, errG\nend\n\n\nfunction BiGANModel:fGx(x, opt)\n self.gradParametersG, self.errG = self:fGx_basic(x, self.netG, self.netD,\n self.gradParametersG, opt)\n return self.errG, self.gradParametersG\nend\n\n\nfunction BiGANModel:fEx_basic(x, netE, netD, gradParametersE, opt)\n util.BiasZero(netE)\n util.BiasZero(netD)\n gradParametersE:zero()\n\n -- First. G(A) should fake the discriminator\n local output = netD:forward(self.fake_AB):clone()\n local errE= self.criterionGAN:forward(output, self.real_label)\n local dgan_loss_dd = self.criterionGAN:backward(output, self.real_label)\n local dgan_loss_do = netD:updateGradInput(self.fake_AB, dgan_loss_dd)\n netE:backward(self.real_B, dgan_loss_do[self.A_idx])-- fake_AB: fake_A, real_B -> fake_label\n return gradParametersE, errE\nend\n\n\nfunction BiGANModel:fEx(x, opt)\n self.gradParametersE, self.errE = self:fEx_basic(x, self.netE, self.netD,\n self.gradParametersE, opt)\n return self.errE, self.gradParametersE\nend\n\n\nfunction BiGANModel:OptimizeParameters(opt)\n local fG = function(x) return self:fGx(x, opt) end\n local fE = function(x) return self:fEx(x, opt) end\n local fD = function(x) return self:fDx(x, opt) end\n optim.adam(fD, self.parametersD, self.optimStateD)\n optim.adam(fG, self.parametersG, self.optimStateG)\n optim.adam(fE, self.parametersE, self.optimStateE)\nend\n\nfunction BiGANModel:RefreshParameters()\n self.parametersD, self.gradParametersD = nil, nil -- nil them to avoid spiking memory\n self.parametersG, self.gradParametersG = nil, nil\n self.parametersE, self.gradParametersE = nil, nil\n -- define parameters of optimization\n self.parametersD, self.gradParametersD = self.netD:getParameters()\n self.parametersG, self.gradParametersG = self.netG:getParameters()\n self.parametersE, self.gradParametersE = self.netE:getParameters()\nend\n\nfunction BiGANModel:Save(prefix, opt)\n util.save_model(self.netG, prefix .. '_net_G.t7', 1)\n util.save_model(self.netE, prefix .. '_net_E.t7', 1)\n util.save_model(self.netD, prefix .. '_net_D.t7', 1)\nend\n\nfunction BiGANModel:GetCurrentErrorDescription()\n description = ('D: %.4f G: %.4f E: %.4f'):format(\n self.errD and self.errD or -1,\n self.errG and self.errG or -1,\n self.errE and self.errE or -1)\n return description\nend\n\nfunction BiGANModel:GetCurrentErrors()\n local errors = {errD=self.errD, errG=self.errG, errE=self.errE}\n return errors\nend\n\n-- returns a string that describes the display plot configuration\nfunction BiGANModel:DisplayPlot(opt)\n return 'errD,errG,errE'\nend\nfunction BiGANModel:UpdateLearningRate(opt)\n local lrd = opt.lr / opt.niter_decay\n local old_lr = self.optimStateD['learningRate']\n local lr = old_lr - lrd\n self.optimStateD['learningRate'] = lr\n self.optimStateG['learningRate'] = lr\n self.optimStateE['learningRate'] = lr\n print(('update learning rate: %f -> %f'):format(old_lr, lr))\nend\n\nlocal function MakeIm3(im)\n -- print('before im_size', im:size())\n local im3 = nil\n if im:size(2) == 1 then\n im3 = torch.repeatTensor(im, 1,3,1,1)\n else\n im3 = im\n end\n -- print('after im_size', im:size())\n -- print('after im3_size', im3:size())\n return im3\nend\nfunction BiGANModel:GetCurrentVisuals(opt, size)\n if not size then\n size = opt.display_winsize\n end\n\n local visuals = {}\n table.insert(visuals, {img=MakeIm3(self.real_A), label='real_A'})\n table.insert(visuals, {img=MakeIm3(self.fake_B), label='fake_B'})\n table.insert(visuals, {img=MakeIm3(self.real_B), label='real_B'})\n table.insert(visuals, {img=MakeIm3(self.fake_A), label='fake_A'})\n return visuals\nend\n" }, { "path": "models/content_gan_model.lua", "content": "local class = require 'class'\nrequire 'models.base_model'\nrequire 'models.architectures'\nrequire 'util.image_pool'\nutil = paths.dofile('../util/util.lua')\ncontent = paths.dofile('../util/content_loss.lua')\n\nContentGANModel = class('ContentGANModel', 'BaseModel')\n\nfunction ContentGANModel:__init(conf)\n BaseModel.__init(self, conf)\n conf = conf or {}\nend\n\nfunction ContentGANModel:model_name()\n return 'ContentGANModel'\nend\n\nfunction ContentGANModel:InitializeStates()\n local optimState = {learningRate=opt.lr, beta1=opt.beta1,}\n return optimState\nend\n-- Defines models and networks\nfunction ContentGANModel:Initialize(opt)\n if opt.test == 0 then\n self.fakePool = ImagePool(opt.pool_size)\n end\n -- define tensors\n self.real_A = torch.Tensor(opt.batchSize, opt.input_nc, opt.fineSize, opt.fineSize)\n self.fake_B = torch.Tensor(opt.batchSize, opt.output_nc, opt.fineSize, opt.fineSize)\n self.real_B = self.fake_B:clone() --torch.Tensor(opt.batchSize, opt.output_nc, opt.fineSize, opt.fineSize)\n\n -- load/define models\n self.criterionGAN = nn.MSECriterion()\n self.criterionContent = nn.AbsCriterion()\n self.contentFunc = content.defineContent(opt.content_loss, opt.layer_name)\n self.netG, self.netD = nil, nil\n if opt.continue_train == 1 then\n if opt.which_epoch then -- which_epoch option exists in test mode\n self.netG = util.load_test_model('G_A', opt)\n self.netD = util.load_test_model('D_A', opt)\n else\n self.netG = util.load_model('G_A', opt)\n self.netD = util.load_model('D_A', opt)\n end\n else\n self.netG = defineG(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG)\n print('netG...', self.netG)\n self.netD = defineD(opt.output_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, false)\n print('netD...', self.netD)\n end\n -- define real/fake labels\n netD_output_size = self.netD:forward(self.real_A):size()\n self.fake_label = torch.Tensor(netD_output_size):fill(0.0)\n self.real_label = torch.Tensor(netD_output_size):fill(1.0) -- no soft smoothing\n self.optimStateD = self:InitializeStates()\n self.optimStateG = self:InitializeStates()\n self:RefreshParameters()\n print('---------- # Learnable Parameters --------------')\n print(('G = %d'):format(self.parametersG:size(1)))\n print(('D = %d'):format(self.parametersD:size(1)))\n print('------------------------------------------------')\n -- os.exit()\nend\n\n-- Runs the forward pass of the network and\n-- saves the result to member variables of the class\nfunction ContentGANModel:Forward(input, opt)\n if opt.which_direction == 'BtoA' then\n local temp = input.real_A\n input.real_A = input.real_B\n input.real_B = temp\n end\n\n self.real_A:copy(input.real_A)\n self.real_B:copy(input.real_B)\n self.fake_B = self.netG:forward(self.real_A):clone()\n -- output = {self.fake_B}\n output = {}\n -- if opt.test == 1 then\n\n -- end\n return output\nend\n\n-- create closure to evaluate f(X) and df/dX of discriminator\nfunction ContentGANModel:fDx_basic(x, gradParams, netD, netG,\n real_target, fake_target, opt)\n util.BiasZero(netD)\n util.BiasZero(netG)\n gradParams:zero()\n\n local errD_real, errD_rec, errD_fake, errD = 0, 0, 0, 0\n -- Real log(D_A(B))\n local output = netD:forward(real_target)\n errD_real = self.criterionGAN:forward(output, self.real_label)\n df_do = self.criterionGAN:backward(output, self.real_label)\n netD:backward(real_target, df_do)\n\n -- Fake + log(1 - D_A(G_A(A)))\n output = netD:forward(fake_target)\n errD_fake = self.criterionGAN:forward(output, self.fake_label)\n df_do = self.criterionGAN:backward(output, self.fake_label)\n netD:backward(fake_target, df_do)\n errD = (errD_real + errD_fake) / 2.0\n -- print('errD', errD\n return errD, gradParams\nend\n\n\nfunction ContentGANModel:fDx(x, opt)\n fake_B = self.fakePool:Query(self.fake_B)\n self.errD, gradParams = self:fDx_basic(x, self.gradparametersD, self.netD, self.netG,\n self.real_B, fake_B, opt)\n return self.errD, gradParams\nend\n\nfunction ContentGANModel:fGx_basic(x, netG_source, netD_source, real_source, real_target, fake_target,\n gradParametersG_source, opt)\n util.BiasZero(netD_source)\n util.BiasZero(netG_source)\n gradParametersG_source:zero()\n -- GAN loss\n -- local df_d_GAN = torch.zeros(fake_target:size())\n -- local errGAN = 0\n -- local errRec = 0\n --- Domain GAN loss: D_A(G_A(A))\n local output = netD_source.output -- [hack] forward was already executed in fDx, so save computation netD_source:forward(fake_B) ---\n local errGAN = self.criterionGAN:forward(output, self.real_label)\n local df_do = self.criterionGAN:backward(output, self.real_label)\n local df_d_GAN = netD_source:updateGradInput(fake_target, df_do) ---:narrow(2,fake_AB:size(2)-output_nc+1, output_nc)\n\n -- content loss\n -- print('content_loss', opt.content_loss)\n -- function content.lossUpdate(criterionContent, real_source, fake_target, contentFunc, loss_type, weight)\n local errContent, df_d_content = content.lossUpdate(self.criterionContent, real_source, fake_target, self.contentFunc, opt.content_loss, opt.lambda_A)\n netG_source:forward(real_source)\n netG_source:backward(real_source, df_d_GAN + df_d_content)\n -- print('errD', errGAN)\n return gradParametersG_source, errGAN, errContent\nend\n\nfunction ContentGANModel:fGx(x, opt)\n self.gradparametersG, self.errG, self.errCont =\n self:fGx_basic(x, self.netG, self.netD,\n self.real_A, self.real_B, self.fake_B,\n self.gradparametersG, opt)\n return self.errG, self.gradparametersG\nend\n\nfunction ContentGANModel:OptimizeParameters(opt)\n local fDx = function(x) return self:fDx(x, opt) end\n local fGx = function(x) return self:fGx(x, opt) end\n optim.adam(fDx, self.parametersD, self.optimStateD)\n optim.adam(fGx, self.parametersG, self.optimStateG)\nend\n\nfunction ContentGANModel:RefreshParameters()\n self.parametersD, self.gradparametersD = nil, nil -- nil them to avoid spiking memory\n self.parametersG, self.gradparametersG = nil, nil\n -- define parameters of optimization\n self.parametersG, self.gradparametersG = self.netG:getParameters()\n self.parametersD, self.gradparametersD = self.netD:getParameters()\nend\n\nfunction ContentGANModel:Save(prefix, opt)\n util.save_model(self.netG, prefix .. '_net_G_A.t7', 1.0)\n util.save_model(self.netD, prefix .. '_net_D_A.t7', 1.0)\nend\n\nfunction ContentGANModel:GetCurrentErrorDescription()\n description = ('G: %.4f D: %.4f Content: %.4f'):format(self.errG and self.errG or -1,\n self.errD and self.errD or -1,\n self.errCont and self.errCont or -1)\n return description\nend\n\n\nfunction ContentGANModel:GetCurrentErrors()\n local errors = {errG=self.errG and self.errG or -1, errD=self.errD and self.errD or -1,\n errCont=self.errCont and self.errCont or -1}\n return errors\nend\n\n-- returns a string that describes the display plot configuration\nfunction ContentGANModel:DisplayPlot(opt)\n return 'errG,errD,errCont'\nend\n\n\nfunction ContentGANModel:GetCurrentVisuals(opt, size)\n if not size then\n size = opt.display_winsize\n end\n\n local visuals = {}\n table.insert(visuals, {img=self.real_A, label='real_A'})\n table.insert(visuals, {img=self.fake_B, label='fake_B'})\n table.insert(visuals, {img=self.real_B, label='real_B'})\n return visuals\nend\n" }, { "path": "models/cycle_gan_model.lua", "content": "local class = require 'class'\nrequire 'models.base_model'\nrequire 'models.architectures'\nrequire 'util.image_pool'\n\nutil = paths.dofile('../util/util.lua')\nCycleGANModel = class('CycleGANModel', 'BaseModel')\n\nfunction CycleGANModel:__init(conf)\n BaseModel.__init(self, conf)\n conf = conf or {}\nend\n\nfunction CycleGANModel:model_name()\n return 'CycleGANModel'\nend\n\nfunction CycleGANModel:InitializeStates(use_wgan)\n optimState = {learningRate=opt.lr, beta1=opt.beta1,}\n return optimState\nend\n-- Defines models and networks\nfunction CycleGANModel:Initialize(opt)\n if opt.test == 0 then\n self.fakeAPool = ImagePool(opt.pool_size)\n self.fakeBPool = ImagePool(opt.pool_size)\n end\n -- define tensors\n if opt.test == 0 then -- allocate tensors for training\n self.real_A = torch.Tensor(opt.batchSize, opt.input_nc, opt.fineSize, opt.fineSize)\n self.real_B = torch.Tensor(opt.batchSize, opt.output_nc, opt.fineSize, opt.fineSize)\n self.fake_A = torch.Tensor(opt.batchSize, opt.input_nc, opt.fineSize, opt.fineSize)\n self.fake_B = torch.Tensor(opt.batchSize, opt.output_nc, opt.fineSize, opt.fineSize)\n self.rec_A = torch.Tensor(opt.batchSize, opt.input_nc, opt.fineSize, opt.fineSize)\n self.rec_B = torch.Tensor(opt.batchSize, opt.output_nc, opt.fineSize, opt.fineSize)\n end\n -- load/define models\n local use_lsgan = ((opt.use_lsgan ~= nil) and (opt.use_lsgan == 1))\n if not use_lsgan then\n self.criterionGAN = nn.BCECriterion()\n else\n self.criterionGAN = nn.MSECriterion()\n end\n self.criterionRec = nn.AbsCriterion()\n\n local netG_A, netD_A, netG_B, netD_B = nil, nil, nil, nil\n if opt.continue_train == 1 then\n if opt.test == 1 then -- test mode\n netG_A = util.load_test_model('G_A', opt)\n netG_B = util.load_test_model('G_B', opt)\n\n --setup optnet to save a little bit of memory\n if opt.use_optnet == 1 then\n local sample_input = torch.randn(1, opt.input_nc, 2, 2)\n local optnet = require 'optnet'\n optnet.optimizeMemory(netG_A, sample_input, {inplace=true, reuseBuffers=true})\n optnet.optimizeMemory(netG_B, sample_input, {inplace=true, reuseBuffers=true})\n end\n else\n netG_A = util.load_model('G_A', opt)\n netG_B = util.load_model('G_B', opt)\n netD_A = util.load_model('D_A', opt)\n netD_B = util.load_model('D_B', opt)\n end\n else\n local use_sigmoid = (not use_lsgan)\n -- netG_test = defineG(opt.input_nc, opt.output_nc, opt.ngf, \"resnet_unet\", opt.arch)\n -- os.exit()\n netG_A = defineG(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.arch)\n print('netG_A...', netG_A)\n netD_A = defineD(opt.output_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, use_sigmoid) -- no sigmoid layer\n print('netD_A...', netD_A)\n netG_B = defineG(opt.output_nc, opt.input_nc, opt.ngf, opt.which_model_netG, opt.arch)\n print('netG_B...', netG_B)\n netD_B = defineD(opt.input_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, use_sigmoid) -- no sigmoid layer\n print('netD_B', netD_B)\n end\n\n self.netD_A = netD_A\n self.netG_A = netG_A\n self.netG_B = netG_B\n self.netD_B = netD_B\n\n -- define real/fake labels\n if opt.test == 0 then\n local D_A_size = self.netD_A:forward(self.real_B):size() -- hack: assume D_size_A = D_size_B\n self.fake_label_A = torch.Tensor(D_A_size):fill(0.0)\n self.real_label_A = torch.Tensor(D_A_size):fill(1.0) -- no soft smoothing\n local D_B_size = self.netD_B:forward(self.real_A):size() -- hack: assume D_size_A = D_size_B\n self.fake_label_B = torch.Tensor(D_B_size):fill(0.0)\n self.real_label_B = torch.Tensor(D_B_size):fill(1.0) -- no soft smoothing\n self.optimStateD_A = self:InitializeStates()\n self.optimStateG_A = self:InitializeStates()\n self.optimStateD_B = self:InitializeStates()\n self.optimStateG_B = self:InitializeStates()\n self:RefreshParameters()\n print('---------- # Learnable Parameters --------------')\n print(('G_A = %d'):format(self.parametersG_A:size(1)))\n print(('D_A = %d'):format(self.parametersD_A:size(1)))\n print(('G_B = %d'):format(self.parametersG_B:size(1)))\n print(('D_B = %d'):format(self.parametersD_B:size(1)))\n print('------------------------------------------------')\n end\nend\n\n-- Runs the forward pass of the network and\n-- saves the result to member variables of the class\nfunction CycleGANModel:Forward(input, opt)\n if opt.which_direction == 'BtoA' then\n \tlocal temp = input.real_A:clone()\n \tinput.real_A = input.real_B:clone()\n \tinput.real_B = temp\n end\n\n if opt.test == 0 then\n self.real_A:copy(input.real_A)\n self.real_B:copy(input.real_B)\n end\n\n if opt.test == 1 then -- forward for test\n if opt.gpu > 0 then\n self.real_A = input.real_A:cuda()\n self.real_B = input.real_B:cuda()\n else\n self.real_A = input.real_A:clone()\n self.real_B = input.real_B:clone()\n end\n self.fake_B = self.netG_A:forward(self.real_A):clone()\n self.fake_A = self.netG_B:forward(self.real_B):clone()\n self.rec_A = self.netG_B:forward(self.fake_B):clone()\n self.rec_B = self.netG_A:forward(self.fake_A):clone()\n end\nend\n\n-- create closure to evaluate f(X) and df/dX of discriminator\nfunction CycleGANModel:fDx_basic(x, gradParams, netD, netG, real, fake, real_label, fake_label, opt)\n util.BiasZero(netD)\n util.BiasZero(netG)\n gradParams:zero()\n -- Real log(D_A(B))\n local output = netD:forward(real)\n local errD_real = self.criterionGAN:forward(output, real_label)\n local df_do = self.criterionGAN:backward(output, real_label)\n netD:backward(real, df_do)\n -- Fake + log(1 - D_A(G_A(A)))\n output = netD:forward(fake)\n local errD_fake = self.criterionGAN:forward(output, fake_label)\n local df_do2 = self.criterionGAN:backward(output, fake_label)\n netD:backward(fake, df_do2)\n -- Compute loss\n local errD = (errD_real + errD_fake) / 2.0\n return errD, gradParams\nend\n\n\nfunction CycleGANModel:fDAx(x, opt)\n -- use image pool that stores the old fake images\n fake_B = self.fakeBPool:Query(self.fake_B)\n self.errD_A, gradParams = self:fDx_basic(x, self.gradparametersD_A, self.netD_A, self.netG_A,\n self.real_B, fake_B, self.real_label_A, self.fake_label_A, opt)\n return self.errD_A, gradParams\nend\n\n\nfunction CycleGANModel:fDBx(x, opt)\n -- use image pool that stores the old fake images\n fake_A = self.fakeAPool:Query(self.fake_A)\n self.errD_B, gradParams = self:fDx_basic(x, self.gradparametersD_B, self.netD_B, self.netG_B,\n self.real_A, fake_A, self.real_label_B, self.fake_label_B, opt)\n return self.errD_B, gradParams\nend\n\n\nfunction CycleGANModel:fGx_basic(x, gradParams, netG, netD, netE, real, real2, real_label, lambda1, lambda2, opt)\n util.BiasZero(netD)\n util.BiasZero(netG)\n util.BiasZero(netE) -- inverse mapping\n gradParams:zero()\n\n -- G should be identity if real2 is fed.\n local errI = nil\n local identity = nil\n if opt.lambda_identity > 0 then\n identity = netG:forward(real2):clone()\n errI = self.criterionRec:forward(identity, real2) * lambda2 * opt.lambda_identity\n local didentity_loss_do = self.criterionRec:backward(identity, real2):mul(lambda2):mul(opt.lambda_identity)\n netG:backward(real2, didentity_loss_do)\n end\n\n --- GAN loss: D_A(G_A(A))\n local fake = netG:forward(real):clone()\n local output = netD:forward(fake)\n local errG = self.criterionGAN:forward(output, real_label)\n local df_do1 = self.criterionGAN:backward(output, real_label)\n local df_d_GAN = netD:updateGradInput(fake, df_do1) --\n\n -- forward cycle loss\n local rec = netE:forward(fake):clone()\n local errRec = self.criterionRec:forward(rec, real) * lambda1\n local df_do2 = self.criterionRec:backward(rec, real):mul(lambda1)\n local df_do_rec = netE:updateGradInput(fake, df_do2)\n\n netG:backward(real, df_d_GAN + df_do_rec)\n\n -- backward cycle loss\n local fake2 = netE:forward(real2)--:clone()\n local rec2 = netG:forward(fake2)--:clone()\n local errAdapt = self.criterionRec:forward(rec2, real2) * lambda2\n local df_do_coadapt = self.criterionRec:backward(rec2, real2):mul(lambda2)\n netG:backward(fake2, df_do_coadapt)\n\n return gradParams, errG, errRec, errI, fake, rec, identity\nend\n\nfunction CycleGANModel:fGAx(x, opt)\n self.gradparametersG_A, self.errG_A, self.errRec_A, self.errI_A, self.fake_B, self.rec_A, self.identity_B =\n self:fGx_basic(x, self.gradparametersG_A, self.netG_A, self.netD_A, self.netG_B, self.real_A, self.real_B,\n self.real_label_A, opt.lambda_A, opt.lambda_B, opt)\n return self.errG_A, self.gradparametersG_A\nend\n\nfunction CycleGANModel:fGBx(x, opt)\n self.gradparametersG_B, self.errG_B, self.errRec_B, self.errI_B, self.fake_A, self.rec_B, self.identity_A =\n self:fGx_basic(x, self.gradparametersG_B, self.netG_B, self.netD_B, self.netG_A, self.real_B, self.real_A,\n self.real_label_B, opt.lambda_B, opt.lambda_A, opt)\n return self.errG_B, self.gradparametersG_B\nend\n\n\nfunction CycleGANModel:OptimizeParameters(opt)\n local fDA = function(x) return self:fDAx(x, opt) end\n local fGA = function(x) return self:fGAx(x, opt) end\n local fDB = function(x) return self:fDBx(x, opt) end\n local fGB = function(x) return self:fGBx(x, opt) end\n\n optim.adam(fGA, self.parametersG_A, self.optimStateG_A)\n optim.adam(fDA, self.parametersD_A, self.optimStateD_A)\n optim.adam(fGB, self.parametersG_B, self.optimStateG_B)\n optim.adam(fDB, self.parametersD_B, self.optimStateD_B)\nend\n\nfunction CycleGANModel:RefreshParameters()\n self.parametersD_A, self.gradparametersD_A = nil, nil -- nil them to avoid spiking memory\n self.parametersG_A, self.gradparametersG_A = nil, nil\n self.parametersG_B, self.gradparametersG_B = nil, nil\n self.parametersD_B, self.gradparametersD_B = nil, nil\n -- define parameters of optimization\n self.parametersG_A, self.gradparametersG_A = self.netG_A:getParameters()\n self.parametersD_A, self.gradparametersD_A = self.netD_A:getParameters()\n self.parametersG_B, self.gradparametersG_B = self.netG_B:getParameters()\n self.parametersD_B, self.gradparametersD_B = self.netD_B:getParameters()\nend\n\nfunction CycleGANModel:Save(prefix, opt)\n util.save_model(self.netG_A, prefix .. '_net_G_A.t7', 1)\n util.save_model(self.netD_A, prefix .. '_net_D_A.t7', 1)\n util.save_model(self.netG_B, prefix .. '_net_G_B.t7', 1)\n util.save_model(self.netD_B, prefix .. '_net_D_B.t7', 1)\nend\n\nfunction CycleGANModel:GetCurrentErrorDescription()\n description = ('[A] G: %.4f D: %.4f Rec: %.4f I: %.4f || [B] G: %.4f D: %.4f Rec: %.4f I:%.4f'):format(\n self.errG_A and self.errG_A or -1,\n self.errD_A and self.errD_A or -1,\n self.errRec_A and self.errRec_A or -1,\n self.errI_A and self.errI_A or -1,\n self.errG_B and self.errG_B or -1,\n self.errD_B and self.errD_B or -1,\n self.errRec_B and self.errRec_B or -1,\n self.errI_B and self.errI_B or -1)\n return description\nend\n\nfunction CycleGANModel:GetCurrentErrors()\n local errors = {errG_A=self.errG_A, errD_A=self.errD_A, errRec_A=self.errRec_A, errI_A=self.errI_A,\n errG_B=self.errG_B, errD_B=self.errD_B, errRec_B=self.errRec_B, errI_B=self.errI_B}\n return errors\nend\n\n-- returns a string that describes the display plot configuration\nfunction CycleGANModel:DisplayPlot(opt)\n if opt.lambda_identity > 0 then\n return 'errG_A,errD_A,errRec_A,errI_A,errG_B,errD_B,errRec_B,errI_B'\n else\n return 'errG_A,errD_A,errRec_A,errG_B,errD_B,errRec_B'\n end\nend\n\nfunction CycleGANModel:UpdateLearningRate(opt)\n local lrd = opt.lr / opt.niter_decay\n local old_lr = self.optimStateD_A['learningRate']\n local lr = old_lr - lrd\n self.optimStateD_A['learningRate'] = lr\n self.optimStateD_B['learningRate'] = lr\n self.optimStateG_A['learningRate'] = lr\n self.optimStateG_B['learningRate'] = lr\n print(('update learning rate: %f -> %f'):format(old_lr, lr))\nend\n\nlocal function MakeIm3(im)\n if im:size(2) == 1 then\n local im3 = torch.repeatTensor(im, 1,3,1,1)\n return im3\n else\n return im\n end\nend\n\nfunction CycleGANModel:GetCurrentVisuals(opt, size)\n local visuals = {}\n table.insert(visuals, {img=MakeIm3(self.real_A), label='real_A'})\n table.insert(visuals, {img=MakeIm3(self.fake_B), label='fake_B'})\n table.insert(visuals, {img=MakeIm3(self.rec_A), label='rec_A'})\n if opt.test == 0 and opt.lambda_identity > 0 then\n table.insert(visuals, {img=MakeIm3(self.identity_A), label='identity_A'})\n end\n table.insert(visuals, {img=MakeIm3(self.real_B), label='real_B'})\n table.insert(visuals, {img=MakeIm3(self.fake_A), label='fake_A'})\n table.insert(visuals, {img=MakeIm3(self.rec_B), label='rec_B'})\n if opt.test == 0 and opt.lambda_identity > 0 then\n table.insert(visuals, {img=MakeIm3(self.identity_B), label='identity_B'})\n end\n return visuals\nend\n" }, { "path": "models/one_direction_test_model.lua", "content": "local class = require 'class'\nrequire 'models.base_model'\nrequire 'models.architectures'\nrequire 'util.image_pool'\n\nutil = paths.dofile('../util/util.lua')\nOneDirectionTestModel = class('OneDirectionTestModel', 'BaseModel')\n\nfunction OneDirectionTestModel:__init(conf)\n BaseModel.__init(self, conf)\n conf = conf or {}\nend\n\nfunction OneDirectionTestModel:model_name()\n return 'OneDirectionTestModel'\nend\n\n-- Defines models and networks\nfunction OneDirectionTestModel:Initialize(opt)\n -- define tensors\n self.real_A = torch.Tensor(opt.batchSize, opt.input_nc, opt.fineSize, opt.fineSize)\n\n -- load/define models\n self.netG_A = util.load_test_model('G', opt)\n\n -- setup optnet to save a bit of memory\n if opt.use_optnet == 1 then\n local optnet = require 'optnet'\n local sample_input = torch.randn(1, opt.input_nc, 2, 2)\n optnet.optimizeMemory(self.netG_A, sample_input, {inplace=true, reuseBuffers=true})\n end\n\n self:RefreshParameters()\n\n print('---------- # Learnable Parameters --------------')\n print(('G_A = %d'):format(self.parametersG_A:size(1)))\n print('------------------------------------------------')\nend\n\n-- Runs the forward pass of the network and\n-- saves the result to member variables of the class\nfunction OneDirectionTestModel:Forward(input, opt)\n if opt.which_direction == 'BtoA' then\n \tinput.real_A = input.real_B:clone()\n end\n\n self.real_A = input.real_A:clone()\n if opt.gpu > 0 then\n self.real_A = self.real_A:cuda()\n end\n\n self.fake_B = self.netG_A:forward(self.real_A):clone()\nend\n\nfunction OneDirectionTestModel:RefreshParameters()\n self.parametersG_A, self.gradparametersG_A = nil, nil\n self.parametersG_A, self.gradparametersG_A = self.netG_A:getParameters()\nend\n\n\nlocal function MakeIm3(im)\n if im:size(2) == 1 then\n local im3 = torch.repeatTensor(im, 1,3,1,1)\n return im3\n else\n return im\n end\nend\n\nfunction OneDirectionTestModel:GetCurrentVisuals(opt, size)\n if not size then\n size = opt.display_winsize\n end\n\n local visuals = {}\n table.insert(visuals, {img=MakeIm3(self.real_A), label='real_A'})\n table.insert(visuals, {img=MakeIm3(self.fake_B), label='fake_B'})\n return visuals\nend\n" }, { "path": "models/pix2pix_model.lua", "content": "local class = require 'class'\nrequire 'models.base_model'\nrequire 'models.architectures'\nrequire 'util.image_pool'\nutil = paths.dofile('../util/util.lua')\nPix2PixModel = class('Pix2PixModel', 'BaseModel')\n\nfunction Pix2PixModel:__init(conf)\n conf = conf or {}\nend\n\n-- Returns the name of the model\nfunction Pix2PixModel:model_name()\n return 'Pix2PixModel'\nend\n\nfunction Pix2PixModel:InitializeStates()\n return {learningRate=opt.lr, beta1=opt.beta1,}\nend\n\n-- Defines models and networks\nfunction Pix2PixModel:Initialize(opt) -- use lsgan\n -- define tensors\n local d_input_nc = opt.input_nc + opt.output_nc\n self.real_AB = torch.Tensor(opt.batchSize, d_input_nc, opt.fineSize, opt.fineSize)\n self.fake_AB = torch.Tensor(opt.batchSize, d_input_nc, opt.fineSize, opt.fineSize)\n if opt.test == 0 then\n self.fakeABPool = ImagePool(opt.pool_size)\n end\n -- load/define models\n self.criterionGAN = nn.MSECriterion()\n self.criterionL1 = nn.AbsCriterion()\n\n local netG, netD = nil, nil\n if opt.continue_train == 1 then\n if opt.test == 1 then -- only load model G for test\n netG = util.load_test_model('G', opt)\n else\n netG = util.load_model('G', opt)\n netD = util.load_model('D', opt)\n end\n else\n netG = defineG(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG)\n netD = defineD(d_input_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, false) -- with sigmoid\n end\n\n self.netD = netD\n self.netG = netG\n\n -- define real/fake labels\n if opt.test == 0 then\n netD_output_size = self.netD:forward(self.real_AB):size()\n self.fake_label = torch.Tensor(netD_output_size):fill(0.0)\n self.real_label = torch.Tensor(netD_output_size):fill(1.0) -- no soft smoothing\n\n self.optimStateD = self:InitializeStates()\n self.optimStateG = self:InitializeStates()\n\n self:RefreshParameters()\n\n print('---------- # Learnable Parameters --------------')\n print(('G = %d'):format(self.parametersG:size(1)))\n print(('D = %d'):format(self.parametersD:size(1)))\n print('------------------------------------------------')\n end\n\n self.A_idx = {{}, {1, opt.input_nc}, {}, {}}\n self.B_idx = {{}, {opt.input_nc+1, opt.input_nc+opt.output_nc}, {}, {}}\nend\n\n-- Runs the forward pass of the network\nfunction Pix2PixModel:Forward(input, opt)\n if opt.which_direction == 'BtoA' then\n \tlocal temp = input.real_A\n \tinput.real_A = input.real_B\n \tinput.real_B = temp\n end\n\n if opt.test == 0 then\n self.real_AB[self.A_idx]:copy(input.real_A)\n self.real_AB[self.B_idx]:copy(input.real_B)\n self.real_A = self.real_AB[self.A_idx]\n self.real_B = self.real_AB[self.B_idx]\n\n self.fake_AB[self.A_idx]:copy(self.real_A)\n self.fake_B = self.netG:forward(self.real_A):clone()\n self.fake_AB[self.B_idx]:copy(self.fake_B)\n else\n if opt.gpu > 0 then\n self.real_A = input.real_A:cuda()\n self.real_B = input.real_B:cuda()\n else\n self.real_A = input.real_A:clone()\n self.real_B = input.real_B:clone()\n end\n self.fake_B = self.netG:forward(self.real_A):clone()\n end\nend\n\n-- create closure to evaluate f(X) and df/dX of discriminator\nfunction Pix2PixModel:fDx_basic(x, gradParams, netD, netG, real, fake, opt)\n util.BiasZero(netD)\n util.BiasZero(netG)\n gradParams:zero()\n\n -- Real log(D(B))\n local output = netD:forward(real)\n local errD_real = self.criterionGAN:forward(output, self.real_label)\n local df_do = self.criterionGAN:backward(output, self.real_label)\n netD:backward(real, df_do)\n -- Fake + log(1 - D(G(A)))\n output = netD:forward(fake)\n local errD_fake = self.criterionGAN:forward(output, self.fake_label)\n local df_do2 = self.criterionGAN:backward(output, self.fake_label)\n netD:backward(fake, df_do2)\n -- calculate loss\n local errD = (errD_real + errD_fake) / 2.0\n return errD, gradParams\nend\n\n\nfunction Pix2PixModel:fDx(x, opt)\n fake_AB = self.fakeABPool:Query(self.fake_AB)\n self.errD, gradParams = self:fDx_basic(x, self.gradParametersD, self.netD, self.netG,\n self.real_AB, fake_AB, opt)\n return self.errD, gradParams\nend\n\nfunction Pix2PixModel:fGx_basic(x, netG, netD, real, fake, gradParametersG, opt)\n util.BiasZero(netG)\n util.BiasZero(netD)\n gradParametersG:zero()\n\n -- First. G(A) should fake the discriminator\n local output = netD:forward(fake)\n local errG = self.criterionGAN:forward(output, self.real_label)\n local dgan_loss_dd = self.criterionGAN:backward(output, self.real_label)\n local dgan_loss_do = netD:updateGradInput(fake, dgan_loss_dd)\n\n -- Second. G(A) should be close to the real\n real_B = real[self.B_idx]\n real_A = real[self.A_idx]\n fake_B = fake[self.B_idx]\n local errL1 = self.criterionL1:forward(fake_B, real_B) * opt.lambda_A\n local dl1_loss_do = self.criterionL1:backward(fake_B, real_B) * opt.lambda_A\n netG:backward(real_A, dgan_loss_do[self.B_idx] + dl1_loss_do)\n\n return gradParametersG, errG, errL1\nend\n\nfunction Pix2PixModel:fGx(x, opt)\n self.gradParametersG, self.errG, self.errL1 = self:fGx_basic(x, self.netG, self.netD,\n self.real_AB, self.fake_AB, self.gradParametersG, opt)\n return self.errG, self.gradParametersG\nend\n\n-- Runs the backprop gradient descent\n-- Corresponds to a single batch of data\nfunction Pix2PixModel:OptimizeParameters(opt)\n local fD = function(x) return self:fDx(x, opt) end\n local fG = function(x) return self:fGx(x, opt) end\n optim.adam(fD, self.parametersD, self.optimStateD)\n optim.adam(fG, self.parametersG, self.optimStateG)\nend\n\n-- This function can be used to reset momentum after each epoch\nfunction Pix2PixModel:RefreshParameters()\n self.parametersD, self.gradParametersD = nil, nil -- nil them to avoid spiking memory\n self.parametersG, self.gradParametersG = nil, nil\n\n -- define parameters of optimization\n self.parametersG, self.gradParametersG = self.netG:getParameters()\n self.parametersD, self.gradParametersD = self.netD:getParameters()\nend\n\n-- This function updates the learning rate; lr for the first opt.niter iterations; graduatlly decreases the lr to 0 for the next opt.niter_decay iterations\nfunction Pix2PixModel:UpdateLearningRate(opt)\n local lrd = opt.lr / opt.niter_decay\n local old_lr = self.optimStateD['learningRate']\n local lr = old_lr - lrd\n self.optimStateD['learningRate'] = lr\n self.optimStateG['learningRate'] = lr\n print(('update learning rate: %f -> %f'):format(old_lr, lr))\nend\n\n\n-- Save the current model to the file system\nfunction Pix2PixModel:Save(prefix, opt)\n util.save_model(self.netG, prefix .. '_net_G.t7', 1.0)\n util.save_model(self.netD, prefix .. '_net_D.t7', 1.0)\nend\n\n-- returns a string that describes the current errors\nfunction Pix2PixModel:GetCurrentErrorDescription()\n description = ('G: %.4f D: %.4f L1: %.4f'):format(\n self.errG and self.errG or -1, self.errD and self.errD or -1, self.errL1 and self.errL1 or -1)\n return description\n\nend\n\n\n-- returns a string that describes the display plot configuration\nfunction Pix2PixModel:DisplayPlot(opt)\n return 'errG,errD,errL1'\nend\n\n\n-- returns current errors\nfunction Pix2PixModel:GetCurrentErrors()\n local errors = {errG=self.errG, errD=self.errD, errL1=self.errL1}\n return errors\nend\n\n-- returns a table of image/label pairs that describe\n-- the current results.\n-- |return|: a table of table. List of image/label pairs\nfunction Pix2PixModel:GetCurrentVisuals(opt, size)\n if not size then\n size = opt.display_winsize\n end\n\n local visuals = {}\n table.insert(visuals, {img=self.real_A, label='real_A'})\n table.insert(visuals, {img=self.fake_B, label='fake_B'})\n table.insert(visuals, {img=self.real_B, label='real_B'})\n\n return visuals\nend\n" }, { "path": "options.lua", "content": "--------------------------------------------------------------------------------\n-- Configure options\n--------------------------------------------------------------------------------\n\nlocal options = {}\n-- options for train\nlocal opt_train = {\n DATA_ROOT = '', -- path to images (should have subfolders 'train', 'val', etc)\n batchSize = 1, -- # images in batch\n loadSize = 143, -- scale images to this size\n fineSize = 128, -- then crop to this size\n ngf = 64, -- # of gen filters in first conv layer\n ndf = 64, -- # of discrim filters in first conv layer\n input_nc = 3, -- # of input image channels\n output_nc = 3, -- # of output image channels\n niter = 100, -- # of iter at starting learning rate\n niter_decay = 100, -- # of iter to linearly decay learning rate to zero\n lr = 0.0002, -- initial learning rate for adam\n beta1 = 0.5, -- momentum term of adam\n ntrain = math.huge, -- # of examples per epoch. math.huge for full dataset\n flip = 1, -- if flip the images for data argumentation\n display_id = 10, -- display window id.\n display_winsize = 128, -- display window size\n display_freq = 25, -- display the current results every display_freq iterations\n gpu = 1, -- gpu = 0 is CPU mode. gpu=X is GPU mode on GPU X\n name = '', -- name of the experiment, should generally be passed on the command line\n which_direction = 'AtoB', -- AtoB or BtoA\n phase = 'train', -- train, val, test, etc\n nThreads = 2, -- # threads for loading data\n save_epoch_freq = 1, -- save a model every save_epoch_freq epochs (does not overwrite previously saved models)\n save_latest_freq = 5000, -- save the latest model every latest_freq sgd iterations (overwrites the previous latest model)\n print_freq = 50, -- print the debug information every print_freq iterations\n save_display_freq = 2500, -- save the current display of results every save_display_freq_iterations\n continue_train = 0, -- if continue training, load the latest model: 1: true, 0: false\n serial_batches = 0, -- if 1, takes images in order to make batches, otherwise takes them randomly\n checkpoints_dir = './checkpoints', -- models are saved here\n cache_dir = './cache', -- cache files are saved here\n cudnn = 1, -- set to 0 to not use cudnn\n which_model_netD = 'basic', -- selects model to use for netD\n which_model_netG = 'resnet_6blocks', -- selects model to use for netG\n norm = 'instance', -- batch or instance normalization\n n_layers_D = 3, -- only used if which_model_netD=='n_layers'\n content_loss = 'pixel', -- content loss type: pixel, vgg\n layer_name = 'pixel', -- layer used in content loss (e.g. relu4_2)\n lambda_A = 10.0, -- weight for cycle loss (A -> B -> A)\n lambda_B = 10.0, -- weight for cycle loss (B -> A -> B)\n model = 'cycle_gan', -- which mode to run. 'cycle_gan', 'pix2pix', 'bigan', 'content_gan'\n use_lsgan = 1, -- if 1, use least square GAN, if 0, use vanilla GAN\n align_data = 0, -- if > 0, use the dataloader for where the images are aligned\n pool_size = 50, -- the size of image buffer that stores previously generated images\n resize_or_crop = 'resize_and_crop', -- resizing/cropping strategy: resize_and_crop | crop | scale_width | scale_height\n lambda_identity = 0.5, -- use identity mapping. Setting opt.lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set opt.lambda_identity = 0.1\n use_optnet = 0, -- use optnet to save GPU memory during test\n}\n\n-- options for test\nlocal opt_test = {\n DATA_ROOT = '', -- path to images (should have subfolders 'train', 'val', etc)\n loadSize = 128, -- scale images to this size\n fineSize = 128, -- then crop to this size\n flip = 0, -- horizontal mirroring data augmentation\n display = 1, -- display samples while training. 0 = false\n display_id = 200, -- display window id.\n gpu = 1, -- gpu = 0 is CPU mode. gpu=X is GPU mode on GPU X\n how_many = 'all', -- how many test images to run (set to all to run on every image found in the data/phase folder)\n phase = 'test', -- train, val, test, etc\n aspect_ratio = 1.0, -- aspect ratio of result images\n norm = 'instance', -- batchnorm or isntance norm\n name = '', -- name of experiment, selects which model to run, should generally should be passed on command line\n input_nc = 3, -- # of input image channels\n output_nc = 3, -- # of output image channels\n serial_batches = 1, -- if 1, takes images in order to make batches, otherwise takes them randomly\n cudnn = 1, -- set to 0 to not use cudnn (untested)\n checkpoints_dir = './checkpoints', -- loads models from here\n cache_dir = './cache', -- cache files are saved here\n results_dir='./results/', -- saves results here\n which_epoch = 'latest', -- which epoch to test? set to 'latest' to use latest cached model\n model = 'cycle_gan', -- which mode to run. 'cycle_gan', 'pix2pix', 'bigan', 'content_gan'; to use pretrained model, select `one_direction_test`\n align_data = 0, -- if > 0, use the dataloader for pix2pix\n which_direction = 'AtoB', -- AtoB or BtoA\n resize_or_crop = 'resize_and_crop', -- resizing/cropping strategy: resize_and_crop | crop | scale_width | scale_height\n}\n\n--------------------------------------------------------------------------------\n-- util functions\n--------------------------------------------------------------------------------\nfunction options.clone(opt)\n local copy = {}\n for orig_key, orig_value in pairs(opt) do\n copy[orig_key] = orig_value\n end\n return copy\nend\n\nfunction options.parse_options(mode)\n if mode == 'train' then\n opt = opt_train\n opt.test = 0\n elseif mode == 'test' then\n opt = opt_test\n opt.test = 1\n else\n print(\"Invalid option [\" .. mode .. \"]\")\n return nil\n end\n\n -- one-line argument parser. parses enviroment variables to override the defaults\n for k,v in pairs(opt) do opt[k] = tonumber(os.getenv(k)) or os.getenv(k) or opt[k] end\n if mode == 'test' then\n opt.nThreads = 1\n opt.continue_train = 1\n opt.batchSize = 1 -- test code only supports batchSize=1\n end\n\n -- print by keys\n keyset = {}\n for k,v in pairs(opt) do\n table.insert(keyset, k)\n end\n table.sort(keyset)\n print(\"------------------- Options -------------------\")\n for i,k in ipairs(keyset) do\n print(('%+25s: %s'):format(k, opt[k]))\n end\n print(\"-----------------------------------------------\")\n\n -- save opt to checkpoints\n paths.mkdir(opt.checkpoints_dir)\n paths.mkdir(paths.concat(opt.checkpoints_dir, opt.name))\n opt.visual_dir = paths.concat(opt.checkpoints_dir, opt.name, 'visuals')\n paths.mkdir(opt.visual_dir)\n -- save opt to the disk\n fd = io.open(paths.concat(opt.checkpoints_dir, opt.name, 'opt_' .. mode .. '.txt'), 'w')\n for i,k in ipairs(keyset) do\n fd:write((\"%+25s: %s\\n\"):format(k, opt[k]))\n end\n fd:close()\n\n return opt\nend\n\n\nreturn options\n" }, { "path": "pretrained_models/download_model.sh", "content": "FILE=$1\n\necho \"Note: available models are apple2orange, facades_photo2label, map2sat, orange2apple, style_cezanne, style_ukiyoe, summer2winter_yosemite, zebra2horse, facades_label2photo, horse2zebra,monet2photo, sat2map, style_monet,style_vangogh, winter2summer_yosemite, iphone2dslr_flower\"\n\necho \"Specified [$FILE]\"\n\nmkdir -p ./checkpoints/${FILE}_pretrained\nURL=https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/models/$FILE.t7\nMODEL_FILE=./checkpoints/${FILE}_pretrained/latest_net_G.t7\nwget -N $URL -O $MODEL_FILE\n" }, { "path": "pretrained_models/download_vgg.sh", "content": "URL1=https://people.eecs.berkeley.edu/~taesung_park/projects/CycleGAN/models/places_vgg.caffemodel\nMODEL_FILE1=./models/places_vgg.caffemodel\nURL2=https://people.eecs.berkeley.edu/~taesung_park/projects/CycleGAN/models/places_vgg.prototxt\nMODEL_FILE2=./models/places_vgg.prototxt\nwget -N $URL1 -O $MODEL_FILE1\nwget -N $URL2 -O $MODEL_FILE2\n" }, { "path": "pretrained_models/places_vgg.prototxt", "content": "name: \"VGG-Places365\"\ninput: \"data\"\ninput_dim: 1\ninput_dim: 3\ninput_dim: 224\ninput_dim: 224\nlayer {\n name: \"conv1_1\"\n type: \"Convolution\"\n bottom: \"data\"\n top: \"conv1_1\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 64\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu1_1\"\n type: \"ReLU\"\n bottom: \"conv1_1\"\n top: \"conv1_1\"\n}\nlayer {\n name: \"conv1_2\"\n type: \"Convolution\"\n bottom: \"conv1_1\"\n top: \"conv1_2\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 64\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu1_2\"\n type: \"ReLU\"\n bottom: \"conv1_2\"\n top: \"conv1_2\"\n}\nlayer {\n name: \"pool1\"\n type: \"Pooling\"\n bottom: \"conv1_2\"\n top: \"pool1\"\n pooling_param {\n pool: MAX\n kernel_size: 2\n stride: 2\n }\n}\nlayer {\n name: \"conv2_1\"\n type: \"Convolution\"\n bottom: \"pool1\"\n top: \"conv2_1\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 128\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu2_1\"\n type: \"ReLU\"\n bottom: \"conv2_1\"\n top: \"conv2_1\"\n}\nlayer {\n name: \"conv2_2\"\n type: \"Convolution\"\n bottom: \"conv2_1\"\n top: \"conv2_2\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 128\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu2_2\"\n type: \"ReLU\"\n bottom: \"conv2_2\"\n top: \"conv2_2\"\n}\nlayer {\n name: \"pool2\"\n type: \"Pooling\"\n bottom: \"conv2_2\"\n top: \"pool2\"\n pooling_param {\n pool: MAX\n kernel_size: 2\n stride: 2\n }\n}\nlayer {\n name: \"conv3_1\"\n type: \"Convolution\"\n bottom: \"pool2\"\n top: \"conv3_1\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 256\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu3_1\"\n type: \"ReLU\"\n bottom: \"conv3_1\"\n top: \"conv3_1\"\n}\nlayer {\n name: \"conv3_2\"\n type: \"Convolution\"\n bottom: \"conv3_1\"\n top: \"conv3_2\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 256\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu3_2\"\n type: \"ReLU\"\n bottom: \"conv3_2\"\n top: \"conv3_2\"\n}\nlayer {\n name: \"conv3_3\"\n type: \"Convolution\"\n bottom: \"conv3_2\"\n top: \"conv3_3\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 256\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu3_3\"\n type: \"ReLU\"\n bottom: \"conv3_3\"\n top: \"conv3_3\"\n}\nlayer {\n name: \"pool3\"\n type: \"Pooling\"\n bottom: \"conv3_3\"\n top: \"pool3\"\n pooling_param {\n pool: MAX\n kernel_size: 2\n stride: 2\n }\n}\nlayer {\n name: \"conv4_1\"\n type: \"Convolution\"\n bottom: \"pool3\"\n top: \"conv4_1\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 512\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu4_1\"\n type: \"ReLU\"\n bottom: \"conv4_1\"\n top: \"conv4_1\"\n}\nlayer {\n name: \"conv4_2\"\n type: \"Convolution\"\n bottom: \"conv4_1\"\n top: \"conv4_2\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 512\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu4_2\"\n type: \"ReLU\"\n bottom: \"conv4_2\"\n top: \"conv4_2\"\n}\nlayer {\n name: \"conv4_3\"\n type: \"Convolution\"\n bottom: \"conv4_2\"\n top: \"conv4_3\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 512\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu4_3\"\n type: \"ReLU\"\n bottom: \"conv4_3\"\n top: \"conv4_3\"\n}\nlayer {\n name: \"pool4\"\n type: \"Pooling\"\n bottom: \"conv4_3\"\n top: \"pool4\"\n pooling_param {\n pool: MAX\n kernel_size: 2\n stride: 2\n }\n}\nlayer {\n name: \"conv5_1\"\n type: \"Convolution\"\n bottom: \"pool4\"\n top: \"conv5_1\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 512\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu5_1\"\n type: \"ReLU\"\n bottom: \"conv5_1\"\n top: \"conv5_1\"\n}\nlayer {\n name: \"conv5_2\"\n type: \"Convolution\"\n bottom: \"conv5_1\"\n top: \"conv5_2\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 512\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu5_2\"\n type: \"ReLU\"\n bottom: \"conv5_2\"\n top: \"conv5_2\"\n}\nlayer {\n name: \"conv5_3\"\n type: \"Convolution\"\n bottom: \"conv5_2\"\n top: \"conv5_3\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n convolution_param {\n num_output: 512\n pad: 1\n kernel_size: 3\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu5_3\"\n type: \"ReLU\"\n bottom: \"conv5_3\"\n top: \"conv5_3\"\n}\nlayer {\n name: \"pool5\"\n type: \"Pooling\"\n bottom: \"conv5_3\"\n top: \"pool5\"\n pooling_param {\n pool: MAX\n kernel_size: 2\n stride: 2\n }\n}\nlayer {\n name: \"fc6\"\n type: \"InnerProduct\"\n bottom: \"pool5\"\n top: \"fc6\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n inner_product_param {\n num_output: 4096\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu6\"\n type: \"ReLU\"\n bottom: \"fc6\"\n top: \"fc6\"\n}\nlayer {\n name: \"drop6\"\n type: \"Dropout\"\n bottom: \"fc6\"\n top: \"fc6\"\n dropout_param {\n dropout_ratio: 0.5\n }\n}\nlayer {\n name: \"fc7\"\n type: \"InnerProduct\"\n bottom: \"fc6\"\n top: \"fc7\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n inner_product_param {\n num_output: 4096\n weight_filler {\n type: \"gaussian\"\n std: 0.01\n }\n bias_filler {\n type: \"constant\"\n value: 0.0\n }\n }\n}\nlayer {\n name: \"relu7\"\n type: \"ReLU\"\n bottom: \"fc7\"\n top: \"fc7\"\n}\nlayer {\n name: \"drop7\"\n type: \"Dropout\"\n bottom: \"fc7\"\n top: \"fc7\"\n dropout_param {\n dropout_ratio: 0.5\n }\n}\nlayer {\n name: \"fc8a\"\n type: \"InnerProduct\"\n bottom: \"fc7\"\n top: \"fc8a\"\n param {\n lr_mult: 1.0\n decay_mult: 1.0\n }\n param {\n lr_mult: 2.0\n decay_mult: 0.0\n }\n inner_product_param {\n num_output: 365\n }\n}\nlayer {\n name: \"prob\"\n type: \"Softmax\"\n bottom: \"fc8a\"\n top: \"prob\"\n}\n" }, { "path": "test.lua", "content": "-- usage: DATA_ROOT=/path/to/data/ name=expt1 which_direction=BtoA th test.lua\n--\n-- code derived from https://github.com/soumith/dcgan.torch and https://github.com/phillipi/pix2pix\nrequire 'image'\nrequire 'nn'\nrequire 'nngraph'\nrequire 'models.architectures'\n\n\nutil = paths.dofile('util/util.lua')\noptions = require 'options'\nopt = options.parse_options('test')\n\n-- initialize torch GPU/CPU mode\nif opt.gpu > 0 then\n require 'cutorch'\n require 'cunn'\n cutorch.setDevice(opt.gpu)\n print (\"GPU Mode\")\n torch.setdefaulttensortype('torch.CudaTensor')\nelse\n torch.setdefaulttensortype('torch.FloatTensor')\n print (\"CPU Mode\")\nend\n\n-- setup visualization\nvisualizer = require 'util/visualizer'\n\nfunction TableConcat(t1,t2)\n for i=1,#t2 do\n t1[#t1+1] = t2[i]\n end\n return t1\nend\n\n\n-- load data\nlocal data_loader = nil\nif opt.align_data > 0 then\n require 'data.aligned_data_loader'\n data_loader = AlignedDataLoader()\nelse\n require 'data.unaligned_data_loader'\n data_loader = UnalignedDataLoader()\nend\nprint( \"DataLoader \" .. data_loader:name() .. \" was created.\")\ndata_loader:Initialize(opt)\n\nif opt.how_many == 'all' then\n opt.how_many = data_loader:size()\nend\n\nopt.how_many = math.min(opt.how_many, data_loader:size())\n\n-- set batch/instance normalization\nset_normalization(opt.norm)\n\n-- load model\nopt.continue_train = 1\n-- define model\nif opt.model == 'cycle_gan' then\n require 'models.cycle_gan_model'\n model = CycleGANModel()\nelseif opt.model == 'one_direction_test' then\n require 'models.one_direction_test_model'\n model = OneDirectionTestModel()\nelseif opt.model == 'pix2pix' then\n require 'models.pix2pix_model'\n model = Pix2PixModel()\nelseif opt.model == 'bigan' then\n require 'models.bigan_model'\n model = BiGANModel()\nelseif opt.model == 'content_gan' then\n require 'models.content_gan_model'\n model = ContentGANModel()\nelse\n error('Please specify a correct model')\nend\nmodel:Initialize(opt)\n\nlocal pathsA = {} -- paths to images A tested on\nlocal pathsB = {} -- paths to images B tested on\nlocal web_dir = paths.concat(opt.results_dir, opt.name .. '/' .. opt.which_epoch .. '_' .. opt.phase)\npaths.mkdir(web_dir)\nlocal image_dir = paths.concat(web_dir, 'images')\npaths.mkdir(image_dir)\ns1 = opt.fineSize\ns2 = opt.fineSize / opt.aspect_ratio\n\nvisuals = {}\n\nfor n = 1, math.floor(opt.how_many) do\n print('processing batch ' .. n)\n local cur_dataA, cur_dataB, cur_pathsA, cur_pathsB = data_loader:GetNextBatch()\n\n cur_pathsA = util.basename_batch(cur_pathsA)\n cur_pathsB = util.basename_batch(cur_pathsB)\n print('pathsA', cur_pathsA)\n print('pathsB', cur_PathsB)\n model:Forward({real_A=cur_dataA, real_B=cur_dataB}, opt)\n\n visuals = model:GetCurrentVisuals(opt, opt.fineSize)\n\n for i,visual in ipairs(visuals) do\n if opt.resize_or_crop == 'scale_width' or opt.resize_or_crop == 'scale_height' then\n s1 = nil\n s2 = nil\n end\n visualizer.save_images(visual.img, paths.concat(image_dir, visual.label), {string.gsub(cur_pathsA[1],'.jpg','.png')}, s1, s2)\n end\n\n\n print('Saved images to: ', image_dir)\n pathsA = TableConcat(pathsA, cur_pathsA)\n pathsB = TableConcat(pathsB, cur_pathsB)\nend\n\nlabels = {}\nfor i,visual in ipairs(visuals) do\n table.insert(labels, visual.label)\nend\n\n-- make webpage\nio.output(paths.concat(web_dir, 'index.html'))\nio.write('
| Image | ')\nfor i = 1, #labels do\n io.write('' .. labels[i] .. ' | ')\nend\nio.write('
| ' .. tostring(n) .. ' | ')\n for j = 1, #labels do\n label = labels[j]\n io.write('![](\"./images/') | ')\n end\n io.write('