[ { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2017 \n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n\n--------------------------- LICENSE FOR mnist-svhn-transfer ---------\n\nMIT License\n\nCopyright (c) 2017\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n\n\n--------------------------- LICENSE FOR pytorch-CycleGAN-and-pix2pix ---------\n\nCopyright (c) 2017, Jun-Yan Zhu and Taesung Park\nAll rights reserved.\n\nRedistribution and use in source and binary forms, with or without\nmodification, are permitted provided that the following conditions are met:\n\n* Redistributions of source code must retain the above copyright notice, this\n list of conditions and the following disclaimer.\n\n* Redistributions in binary form must reproduce the above copyright notice,\n this list of conditions and the following disclaimer in the documentation\n and/or other materials provided with the distribution.\n\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\nAND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\nIMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\nDISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\nFOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\nDAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\nSERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\nCAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\nOR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\nOF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n" }, { "path": "README.md", "content": "# Pytorch implementation of One-Shot Unsupervised Cross Domain Translation ([arxiv](https://arxiv.org/abs/1806.06029)).\n\nPrerequisites\n--------------\n- Python 3.6\n- Pytorch 0.4\n- Numpy/Scipy/Pandas\n- Progressbar\n- OpenCV\n- [visdom](https://github.com/facebookresearch/visdom)\n- [dominate](https://github.com/Knio/dominate)\n\n## MNIST-to-SVHN and SVHN-to-MNIST\n\nTo train autoencoder for both MNIST and SVHN (In mnist_to_svhn folder):\npython main_autoencoder.py --use_augmentation=True\n\nTo train OST for MNIST to SVHN:\npython main_mnist_to_svhn.py --pretrained_g=True --save_models_and_samples=True --use_augmentation=True --one_way_cycle=True --freeze_shared=False\n\nTo train OST for SVHN to MNIST:\npython main_svhn_to_mnist.py --pretrained_g=True --save_models_and_samples=True --use_augmentation=True --one_way_cycle=True --freeze_shared=False\n\n## Drawing and Style Transfer Tasks\n\n### Download Dataset\n\nTo download dataset (in drawing_and_style_transfer folder):\nbash datasets/download_cyclegan_dataset.sh $DATASET_NAME\nwhere DATASET_NAME is one of (facades, cityscapes, maps, monet2photo, summer2winter_yosemite)\n\n### Train Autoencoder\n\nTo train autoencoder for facades (in drawing_and_style_transfer folder):\npython train.py --dataroot=./datasets/facades/trainB --name=facades_autoencoder --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=2\n\nIn the reverse direction (images of facades):\npython train.py --dataroot=./datasets/facades/trainA --name=facades_autoencoder_reverse --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=2\n\n### Train OST\n\nTo train OST for images to facades:\npython train.py --dataroot=./datasets/facades/ --name=facades_ost --load_dir=facades_autoencoder --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1\n\nTo train OST for facades to images (reverse direction):\npython train.py --dataroot=./datasets/facades/ --name=facades_ost_reverse --load_dir=facades_autoencoder_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1 --A='B' --B='A'\n\nTo visualize losses: run python -m visdom.server\n\n### Test OST\n\nTo test OST for images to facades:\npython test.py --dataroot=./datasets/facades/ --name=facades_ost --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1\n\nTo test OST for facades to images (reverse direction):\npython test.py --dataroot=./datasets/facades/ --name=facades_ost_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1 --A='B' --B='A'\n\n### Options\nAdditional scripts for other datasets are at ./drawing_and_style_transfer/scripts\n\nOptions are at ./drawing_and_style_transfer/options\n\n## Reference\nIf you found this code useful, please cite the following paper:\n```\n@inproceedings{Benaim2018OneShotUC,\n title={One-Shot Unsupervised Cross Domain Translation},\n author={Sagie Benaim and Lior Wolf},\n booktitle={NeurIPS},\n year={2018}\n}\n```\n\n" }, { "path": "drawing_and_style_transfer/data/__init__.py", "content": "import torch.utils.data\nfrom data.base_data_loader import BaseDataLoader\n\n\ndef CreateDataLoader(opt):\n data_loader = CustomDatasetDataLoader()\n print(data_loader.name())\n data_loader.initialize(opt)\n return data_loader\n\n\ndef CreateDataset(opt):\n if opt.dataset_mode == 'aligned':\n from data.aligned_dataset import AlignedDataset\n dataset = AlignedDataset()\n elif opt.dataset_mode == 'unaligned':\n from data.unaligned_dataset import UnalignedDataset\n dataset = UnalignedDataset()\n elif opt.dataset_mode == 'single':\n from data.single_dataset import SingleDataset\n dataset = SingleDataset()\n else:\n raise ValueError(\"Dataset [%s] not recognized.\" % opt.dataset_mode)\n\n print(\"dataset [%s] was created\" % (dataset.name()))\n dataset.initialize(opt)\n return dataset\n\n\nclass CustomDatasetDataLoader(BaseDataLoader):\n def name(self):\n return 'CustomDatasetDataLoader'\n\n def initialize(self, opt):\n BaseDataLoader.initialize(self, opt)\n self.dataset = CreateDataset(opt)\n self.dataloader = torch.utils.data.DataLoader(\n self.dataset,\n batch_size=opt.batchSize,\n shuffle=not opt.serial_batches,\n num_workers=int(opt.nThreads))\n\n def load_data(self):\n return self\n\n def __len__(self):\n return len(self.dataset)\n\n def __iter__(self):\n for i, data in enumerate(self.dataloader):\n yield data\n" }, { "path": "drawing_and_style_transfer/data/aligned_dataset.py", "content": "import os.path\nimport random\nimport torchvision.transforms as transforms\nimport torch\nfrom data.base_dataset import BaseDataset\nfrom data.image_folder import make_dataset\nfrom PIL import Image\n\n\nclass AlignedDataset(BaseDataset):\n def initialize(self, opt):\n self.opt = opt\n self.root = opt.dataroot\n self.dir_AB = os.path.join(opt.dataroot, opt.phase)\n self.AB_paths = sorted(make_dataset(self.dir_AB))\n assert (opt.resize_or_crop == 'resize_and_crop')\n\n def __getitem__(self, index):\n AB_path = self.AB_paths[index]\n AB = Image.open(AB_path).convert('RGB')\n w, h = AB.size\n w2 = int(w / 2)\n A = AB.crop((0, 0, w2, h)).resize((self.opt.loadSize, self.opt.loadSize), Image.BICUBIC)\n B = AB.crop((w2, 0, w, h)).resize((self.opt.loadSize, self.opt.loadSize), Image.BICUBIC)\n A = transforms.ToTensor()(A)\n B = transforms.ToTensor()(B)\n w_offset = random.randint(0, max(0, self.opt.loadSize - self.opt.fineSize - 1))\n h_offset = random.randint(0, max(0, self.opt.loadSize - self.opt.fineSize - 1))\n\n A = A[:, h_offset:h_offset + self.opt.fineSize, w_offset:w_offset + self.opt.fineSize]\n B = B[:, h_offset:h_offset + self.opt.fineSize, w_offset:w_offset + self.opt.fineSize]\n\n A = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(A)\n B = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(B)\n\n if self.opt.which_direction == 'BtoA':\n input_nc = self.opt.output_nc\n output_nc = self.opt.input_nc\n else:\n input_nc = self.opt.input_nc\n output_nc = self.opt.output_nc\n\n if (not self.opt.no_flip) and random.random() < 0.5:\n idx = [i for i in range(A.size(2) - 1, -1, -1)]\n idx = torch.LongTensor(idx)\n A = A.index_select(2, idx)\n B = B.index_select(2, idx)\n\n if input_nc == 1: # RGB to gray\n tmp = A[0, ...] * 0.299 + A[1, ...] * 0.587 + A[2, ...] * 0.114\n A = tmp.unsqueeze(0)\n\n if output_nc == 1: # RGB to gray\n tmp = B[0, ...] * 0.299 + B[1, ...] * 0.587 + B[2, ...] * 0.114\n B = tmp.unsqueeze(0)\n\n return {'A': A, 'B': B,\n 'A_paths': AB_path, 'B_paths': AB_path}\n\n def __len__(self):\n return len(self.AB_paths)\n\n def name(self):\n return 'AlignedDataset'\n" }, { "path": "drawing_and_style_transfer/data/base_data_loader.py", "content": "class BaseDataLoader():\n def __init__(self):\n pass\n\n def initialize(self, opt):\n self.opt = opt\n pass\n\n def load_data(self):\n return None\n" }, { "path": "drawing_and_style_transfer/data/base_dataset.py", "content": "import torch.utils.data as data\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n\nclass BaseDataset(data.Dataset):\n def __init__(self):\n super(BaseDataset, self).__init__()\n\n def name(self):\n return 'BaseDataset'\n\n def initialize(self, opt):\n pass\n\n\ndef get_transform(opt):\n transform_list = []\n if opt.resize_or_crop == 'resize_and_crop':\n osize = [opt.loadSize, opt.loadSize]\n transform_list.append(transforms.Scale(osize, Image.BICUBIC))\n transform_list.append(transforms.RandomCrop(opt.fineSize))\n elif opt.resize_or_crop == 'crop':\n transform_list.append(transforms.RandomCrop(opt.fineSize))\n elif opt.resize_or_crop == 'scale_width':\n transform_list.append(transforms.Lambda(\n lambda img: __scale_width(img, opt.fineSize)))\n elif opt.resize_or_crop == 'scale_width_and_crop':\n transform_list.append(transforms.Lambda(\n lambda img: __scale_width(img, opt.loadSize)))\n transform_list.append(transforms.RandomCrop(opt.fineSize))\n\n if opt.isTrain and not opt.no_flip_and_rotation:\n # Default augmentations as in paper\n transform_list.append(transforms.RandomHorizontalFlip())\n transform_list.append(transforms.RandomRotation(opt.rotation_degree))\n\n transform_list += [transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5),\n (0.5, 0.5, 0.5))]\n return transforms.Compose(transform_list)\n\n\ndef __scale_width(img, target_width):\n ow, oh = img.size\n if (ow == target_width):\n return img\n w = target_width\n h = int(target_width * oh / ow)\n return img.resize((w, h), Image.BICUBIC)\n" }, { "path": "drawing_and_style_transfer/data/image_folder.py", "content": "###############################################################################\n# Code from\n# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py\n# Modified the original code so that it also loads images from the current\n# directory as well as the subdirectories\n###############################################################################\n\nimport torch.utils.data as data\n\nfrom PIL import Image\nimport os\nimport os.path\n\nIMG_EXTENSIONS = [\n '.jpg', '.JPG', '.jpeg', '.JPEG',\n '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',\n]\n\n\ndef is_image_file(filename):\n return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)\n\n\ndef make_dataset(dir, max_items=-1, start=0):\n images = []\n assert os.path.isdir(dir), '%s is not a valid directory' % dir\n\n for root, _, fnames in sorted(os.walk(dir)):\n for fname in fnames:\n if is_image_file(fname):\n path = os.path.join(root, fname)\n images.append(path)\n\n if max_items >= 0:\n return sorted(images)[start:start + max_items]\n return images\n\n\ndef default_loader(path):\n return Image.open(path).convert('RGB')\n\n\nclass ImageFolder(data.Dataset):\n def __init__(self, root, transform=None, return_paths=False,\n loader=default_loader):\n imgs = make_dataset(root)\n if len(imgs) == 0:\n raise (RuntimeError(\"Found 0 images in: \" + root + \"\\n\"\n \"Supported image extensions are: \" +\n \",\".join(IMG_EXTENSIONS)))\n\n self.root = root\n self.imgs = imgs\n self.transform = transform\n self.return_paths = return_paths\n self.loader = loader\n\n def __getitem__(self, index):\n path = self.imgs[index]\n img = self.loader(path)\n if self.transform is not None:\n img = self.transform(img)\n if self.return_paths:\n return img, path\n else:\n return img\n\n def __len__(self):\n return len(self.imgs)\n" }, { "path": "drawing_and_style_transfer/data/single_dataset.py", "content": "import os.path\nfrom data.base_dataset import BaseDataset, get_transform\nfrom data.image_folder import make_dataset\nfrom PIL import Image\n\n\nclass SingleDataset(BaseDataset):\n def initialize(self, opt):\n self.opt = opt\n self.root = opt.dataroot\n self.dir_A = os.path.join(opt.dataroot)\n\n self.A_paths = make_dataset(self.dir_A)\n\n self.A_paths = sorted(self.A_paths)\n\n self.transform = get_transform(opt)\n\n def __getitem__(self, index):\n A_path = self.A_paths[index]\n A_img = Image.open(A_path).convert('RGB')\n A = self.transform(A_img)\n if self.opt.which_direction == 'BtoA':\n input_nc = self.opt.output_nc\n else:\n input_nc = self.opt.input_nc\n\n if input_nc == 1: # RGB to gray\n tmp = A[0, ...] * 0.299 + A[1, ...] * 0.587 + A[2, ...] * 0.114\n A = tmp.unsqueeze(0)\n\n return {'A': A, 'A_paths': A_path}\n\n def __len__(self):\n return len(self.A_paths)\n\n def name(self):\n return 'SingleImageDataset'\n" }, { "path": "drawing_and_style_transfer/data/unaligned_dataset.py", "content": "import os.path\nfrom data.base_dataset import BaseDataset, get_transform\nfrom data.image_folder import make_dataset\nfrom PIL import Image\nimport random\n\n\nclass UnalignedDataset(BaseDataset):\n def initialize(self, opt):\n self.opt = opt\n self.root = opt.dataroot\n self.dir_A = os.path.join(opt.dataroot, opt.phase + opt.A)\n self.dir_B = os.path.join(opt.dataroot, opt.phase + opt.B)\n self.A_paths = make_dataset(self.dir_A, max_items=opt.max_items_A, start=opt.start)\n self.B_paths = make_dataset(self.dir_B, max_items=opt.max_items_B, start=opt.start)\n\n self.A_paths = sorted(self.A_paths)\n self.B_paths = sorted(self.B_paths)\n self.A_size = len(self.A_paths)\n self.B_size = len(self.B_paths)\n self.transform = get_transform(opt)\n\n def __getitem__(self, index):\n A_path = self.A_paths[index % self.A_size]\n if self.opt.serial_batches:\n index_B = index % self.B_size\n else:\n index_B = random.randint(0, self.B_size - 1)\n B_path = self.B_paths[index_B]\n # print('(A, B) = (%d, %d)' % (index_A, index_B))\n A_img = Image.open(A_path).convert('RGB')\n B_img = Image.open(B_path).convert('RGB')\n\n A = self.transform(A_img)\n B = self.transform(B_img)\n if self.opt.which_direction == 'BtoA':\n input_nc = self.opt.output_nc\n output_nc = self.opt.input_nc\n else:\n input_nc = self.opt.input_nc\n output_nc = self.opt.output_nc\n\n if input_nc == 1: # RGB to gray\n tmp = A[0, ...] * 0.299 + A[1, ...] * 0.587 + A[2, ...] * 0.114\n A = tmp.unsqueeze(0)\n\n if output_nc == 1: # RGB to gray\n tmp = B[0, ...] * 0.299 + B[1, ...] * 0.587 + B[2, ...] * 0.114\n B = tmp.unsqueeze(0)\n return {'A': A, 'B': B,\n 'A_paths': A_path, 'B_paths': B_path}\n\n def __len__(self):\n return max(self.A_size, self.B_size)\n\n def name(self):\n return 'UnalignedDataset'\n" }, { "path": "drawing_and_style_transfer/datasets/combine_A_and_B.py", "content": "import os\nimport numpy as np\nimport cv2\nimport argparse\n\nparser = argparse.ArgumentParser('create image pairs')\nparser.add_argument('--fold_A', dest='fold_A', help='input directory for image A', type=str, default='../dataset/50kshoes_edges')\nparser.add_argument('--fold_B', dest='fold_B', help='input directory for image B', type=str, default='../dataset/50kshoes_jpg')\nparser.add_argument('--fold_AB', dest='fold_AB', help='output directory', type=str, default='../dataset/test_AB')\nparser.add_argument('--num_imgs', dest='num_imgs', help='number of images',type=int, default=1000000)\nparser.add_argument('--use_AB', dest='use_AB', help='if true: (0001_A, 0001_B) to (0001_AB)',action='store_true')\nargs = parser.parse_args()\n\nfor arg in vars(args):\n print('[%s] = ' % arg, getattr(args, arg))\n\nsplits = os.listdir(args.fold_A)\n\nfor sp in splits:\n img_fold_A = os.path.join(args.fold_A, sp)\n img_fold_B = os.path.join(args.fold_B, sp)\n img_list = os.listdir(img_fold_A)\n if args.use_AB:\n img_list = [img_path for img_path in img_list if '_A.' in img_path]\n\n num_imgs = min(args.num_imgs, len(img_list))\n print('split = %s, use %d/%d images' % (sp, num_imgs, len(img_list)))\n img_fold_AB = os.path.join(args.fold_AB, sp)\n if not os.path.isdir(img_fold_AB):\n os.makedirs(img_fold_AB)\n print('split = %s, number of images = %d' % (sp, num_imgs))\n for n in range(num_imgs):\n name_A = img_list[n]\n path_A = os.path.join(img_fold_A, name_A)\n if args.use_AB:\n name_B = name_A.replace('_A.', '_B.')\n else:\n name_B = name_A\n path_B = os.path.join(img_fold_B, name_B)\n if os.path.isfile(path_A) and os.path.isfile(path_B):\n name_AB = name_A\n if args.use_AB:\n name_AB = name_AB.replace('_A.', '.') # remove _A\n path_AB = os.path.join(img_fold_AB, name_AB)\n im_A = cv2.imread(path_A, cv2.CV_LOAD_IMAGE_COLOR)\n im_B = cv2.imread(path_B, cv2.CV_LOAD_IMAGE_COLOR)\n im_AB = np.concatenate([im_A, im_B], 1)\n cv2.imwrite(path_AB, im_AB)\n" }, { "path": "drawing_and_style_transfer/datasets/download_cyclegan_dataset.sh", "content": "FILE=$1\n\nif [[ $FILE != \"ae_photos\" && $FILE != \"apple2orange\" && $FILE != \"summer2winter_yosemite\" && $FILE != \"horse2zebra\" && $FILE != \"monet2photo\" && $FILE != \"cezanne2photo\" && $FILE != \"ukiyoe2photo\" && $FILE != \"vangogh2photo\" && $FILE != \"maps\" && $FILE != \"cityscapes\" && $FILE != \"facades\" && $FILE != \"iphone2dslr_flower\" && $FILE != \"ae_photos\" ]]; then\n echo \"Available datasets are: apple2orange, summer2winter_yosemite, horse2zebra, monet2photo, cezanne2photo, ukiyoe2photo, vangogh2photo, maps, cityscapes, facades, iphone2dslr_flower, ae_photos\"\n exit 1\nfi\n\nURL=https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets/$FILE.zip\nZIP_FILE=./datasets/$FILE.zip\nTARGET_DIR=./datasets/$FILE/\nwget -N $URL -O $ZIP_FILE\nmkdir $TARGET_DIR\nunzip $ZIP_FILE -d ./datasets/\nrm $ZIP_FILE\n" }, { "path": "drawing_and_style_transfer/datasets/make_dataset_aligned.py", "content": "import os\n\nfrom PIL import Image\n\n\ndef get_file_paths(folder):\n image_file_paths = []\n for root, dirs, filenames in os.walk(folder):\n filenames = sorted(filenames)\n for filename in filenames:\n input_path = os.path.abspath(root)\n file_path = os.path.join(input_path, filename)\n if filename.endswith('.png') or filename.endswith('.jpg'):\n image_file_paths.append(file_path)\n\n break # prevent descending into subfolders\n return image_file_paths\n\n\ndef align_images(a_file_paths, b_file_paths, target_path):\n if not os.path.exists(target_path):\n os.makedirs(target_path)\n\n for i in range(len(a_file_paths)):\n img_a = Image.open(a_file_paths[i])\n img_b = Image.open(b_file_paths[i])\n assert(img_a.size == img_b.size)\n\n aligned_image = Image.new(\"RGB\", (img_a.size[0] * 2, img_a.size[1]))\n aligned_image.paste(img_a, (0, 0))\n aligned_image.paste(img_b, (img_a.size[0], 0))\n aligned_image.save(os.path.join(target_path, '{:04d}.jpg'.format(i)))\n\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser()\n parser.add_argument(\n '--dataset-path',\n dest='dataset_path',\n help='Which folder to process (it should have subfolders testA, testB, trainA and trainB'\n )\n args = parser.parse_args()\n\n dataset_folder = args.dataset_path\n print(dataset_folder)\n\n test_a_path = os.path.join(dataset_folder, 'testA')\n test_b_path = os.path.join(dataset_folder, 'testB')\n test_a_file_paths = get_file_paths(test_a_path)\n test_b_file_paths = get_file_paths(test_b_path)\n assert(len(test_a_file_paths) == len(test_b_file_paths))\n test_path = os.path.join(dataset_folder, 'test')\n\n train_a_path = os.path.join(dataset_folder, 'trainA')\n train_b_path = os.path.join(dataset_folder, 'trainB')\n train_a_file_paths = get_file_paths(train_a_path)\n train_b_file_paths = get_file_paths(train_b_path)\n assert(len(train_a_file_paths) == len(train_b_file_paths))\n train_path = os.path.join(dataset_folder, 'train')\n\n align_images(test_a_file_paths, test_b_file_paths, test_path)\n align_images(train_a_file_paths, train_b_file_paths, train_path)\n" }, { "path": "drawing_and_style_transfer/environment.yml", "content": "name: OST\nchannels:\n- peterjc123\n- defaults\ndependencies:\n- python=3.6.5\n- pytorch=0.4.0\n- scipy\n- pip:\n - dominate==2.3.1\n - git+https://github.com/pytorch/vision.git\n - Pillow==5.0.0\n - numpy==1.14.1\n - visdom==0.1.7\n" }, { "path": "drawing_and_style_transfer/models/__init__.py", "content": "def create_model(opt):\n print(opt.model)\n if opt.model == 'ost':\n assert (opt.dataset_mode == 'unaligned')\n from .ost import OSTModel\n model = OSTModel()\n elif opt.model == 'autoencoder':\n assert (opt.dataset_mode == 'single')\n from .autoencoder_model import AutoEncoderModel\n model = AutoEncoderModel()\n elif opt.model == 'test':\n assert (opt.dataset_mode == 'single')\n from .test_model import TestModel\n model = TestModel()\n else:\n raise NotImplementedError('model [%s] not implemented.' % opt.model)\n model.initialize(opt)\n print(\"model [%s] was created\" % (model.name()))\n return model\n" }, { "path": "drawing_and_style_transfer/models/autoencoder_model.py", "content": "import torch\nfrom collections import OrderedDict\nfrom torch.autograd import Variable\nimport itertools\nimport util.util as util\nfrom util.image_pool import ImagePool\nfrom .base_model import BaseModel\nfrom . import networks\n\n\nclass AutoEncoderModel(BaseModel):\n def name(self):\n return 'AutoEncoderModel'\n\n def set_encoders_and_decoders(self, opt):\n n_downsampling = opt.n_downsampling\n start_unshared = 0\n num_unshared = opt.num_unshared\n start_shared = num_unshared\n end_shared = n_downsampling\n start_dec_shared = start_unshared\n end_dec_shared = start_unshared + (end_shared - start_shared)\n start_dec_unshared = end_dec_shared\n end_dec_unshared = n_downsampling\n\n num_res_blocks_unshared = opt.num_res_blocks_unshared\n n_res_blocks_shared = opt.num_res_blocks_shared\n\n self.netEnc_b, self.netDec_b = networks.define_ED(opt.input_nc, opt.output_nc,\n opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout,\n opt.init_type, self.gpu_ids,\n n_blocks_encoder=num_res_blocks_unshared,\n n_blocks_decoder=num_res_blocks_unshared,\n start=start_unshared,\n end=num_unshared, n_downsampling=n_downsampling,\n input_layer=True,\n output_layer=True, start_dec=start_dec_unshared,\n end_dec=end_dec_unshared)\n\n self.netEnc_shared, self.netDec_shared = networks.define_ED(opt.input_nc, opt.output_nc,\n opt.ngf, opt.which_model_netG, opt.norm,\n not opt.no_dropout,\n opt.init_type, self.gpu_ids,\n n_blocks_encoder=n_res_blocks_shared,\n n_blocks_decoder=n_res_blocks_shared,\n start=start_shared, n_downsampling=n_downsampling,\n end=end_shared,\n input_layer=False,\n output_layer=False, start_dec=start_dec_shared,\n end_dec=end_dec_shared)\n\n def initialize(self, opt):\n BaseModel.initialize(self, opt)\n self.set_encoders_and_decoders(opt)\n\n if self.isTrain:\n use_sigmoid = opt.no_lsgan\n self.netD = networks.define_D(opt.output_nc, opt.ndf,\n opt.which_model_netD,\n opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, self.gpu_ids)\n\n if not self.isTrain or opt.continue_train:\n which_epoch = opt.which_epoch\n self.load_network(self.netEnc_b, 'Enc_b', which_epoch)\n self.load_network(self.netDec_b, 'Dec_b', which_epoch)\n self.load_network(self.netEnc_shared, 'Enc_shared', which_epoch)\n self.load_network(self.netDec_shared, 'Dec_shared', which_epoch)\n if self.isTrain:\n self.load_network(self.netD, 'D', which_epoch)\n\n if self.isTrain:\n self.fake_B_pool = ImagePool(opt.pool_size)\n # define loss functions\n self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)\n self.criterionIdt = torch.nn.L1Loss()\n\n # initialize optimizers\n self.optimizer_Enc = torch.optim.Adam(\n itertools.chain(self.netEnc_b.parameters(), self.netEnc_shared.parameters()),\n lr=opt.lr, betas=(opt.beta1, 0.999))\n self.optimizer_Dec = torch.optim.Adam(\n itertools.chain(self.netDec_b.parameters(), self.netDec_shared.parameters()),\n lr=opt.lr, betas=(opt.beta1, 0.999))\n self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))\n\n self.optimizers = []\n self.schedulers = []\n self.optimizers.append(self.optimizer_Enc)\n self.optimizers.append(self.optimizer_Dec)\n self.optimizers.append(self.optimizer_D)\n for optimizer in self.optimizers:\n self.schedulers.append(networks.get_scheduler(optimizer, opt))\n\n print('---------- Networks initialized -------------')\n networks.print_network(self.netEnc_b)\n networks.print_network(self.netDec_b)\n networks.print_network(self.netEnc_shared)\n networks.print_network(self.netDec_shared)\n if self.isTrain:\n networks.print_network(self.netD)\n print('-----------------------------------------------')\n\n def set_input(self, input):\n # 'A' is given as single_dataset\n input_B = input['A']\n if len(self.gpu_ids) > 0:\n input_B = input_B.cuda(self.gpu_ids[0], async=True)\n self.input_B = input_B\n # 'A' is given as single_dataset\n self.image_paths = input['A_paths']\n\n def forward(self):\n self.real_B = Variable(self.input_B)\n\n def netEnc(self, x):\n return self.netEnc_shared(self.netEnc_b(x))\n\n def netDec(self, x):\n return self.netDec_b(self.netDec_shared(x))\n\n def test(self):\n real_B = Variable(self.input_B, volatile=True)\n fake_B = self.netDec(self.netEnc(real_B))\n self.fake_B = fake_B.data\n\n # get image paths\n def get_image_paths(self):\n return self.image_paths\n\n def backward_D_basic(self, netD, real, fake):\n # Real\n pred_real = netD(real)\n loss_D_real = self.criterionGAN(pred_real, True)\n # Fake\n pred_fake = netD(fake.detach())\n loss_D_fake = self.criterionGAN(pred_fake, False)\n # Combined loss\n loss_D = (loss_D_real + loss_D_fake) * 0.5\n # backward\n loss_D.backward()\n return loss_D\n\n def backward_D(self):\n fake_B = self.fake_B_pool.query(self.fake_B)\n loss_D = self.backward_D_basic(self.netD, self.real_B, fake_B)\n self.loss_D = loss_D.data[0]\n\n def _compute_kl(self, mu):\n mu_2 = torch.pow(mu, 2)\n encoding_loss = torch.mean(mu_2)\n return encoding_loss\n\n def backward_G(self):\n lambda_B = self.opt.lambda_B\n\n # GAN loss D_B(G_B(B))\n enc_b = self.netEnc(self.real_B)\n fake_B = self.netDec(enc_b)\n pred_fake = self.netD(fake_B)\n loss_Gan = self.criterionGAN(pred_fake, True)\n loss_idt_B = self.criterionIdt(fake_B, self.real_B) * lambda_B\n loss_kl_B = self.opt.kl_lambda * self._compute_kl(enc_b)\n\n # combined loss\n loss_G = loss_Gan + loss_idt_B + loss_kl_B\n loss_G.backward()\n\n self.fake_B = fake_B.data\n self.loss_Gan = loss_Gan.data[0]\n self.loss_idt_B = loss_idt_B.data[0]\n\n def optimize_parameters(self):\n # forward\n self.forward()\n\n # G\n self.optimizer_Enc.zero_grad()\n self.optimizer_Dec.zero_grad()\n self.backward_G()\n self.optimizer_Enc.step()\n self.optimizer_Dec.step()\n\n # D\n self.optimizer_D.zero_grad()\n self.backward_D()\n self.optimizer_D.step()\n\n def get_current_errors(self):\n ret_errors = OrderedDict([('D', self.loss_D), ('G_B', self.loss_Gan), ('Idt_B', self.loss_idt_B)])\n return ret_errors\n\n def get_current_visuals(self):\n real_B = util.tensor2im(self.input_B)\n fake_B = util.tensor2im(self.fake_B)\n ret_visuals = OrderedDict([('real_B', real_B), ('fake_B', fake_B), ])\n return ret_visuals\n\n def save(self, label):\n self.save_network(self.netEnc_b, 'Enc_b', label, self.gpu_ids)\n self.save_network(self.netDec_b, 'Dec_b', label, self.gpu_ids)\n self.save_network(self.netEnc_shared, 'Enc_shared', label, self.gpu_ids)\n self.save_network(self.netDec_shared, 'Dec_shared', label, self.gpu_ids)\n self.save_network(self.netD, 'D', label, self.gpu_ids)\n" }, { "path": "drawing_and_style_transfer/models/base_model.py", "content": "import os\nimport torch\n\n\nclass BaseModel(object):\n def name(self):\n return 'BaseModel'\n\n def initialize(self, opt):\n self.opt = opt\n self.gpu_ids = opt.gpu_ids\n self.isTrain = opt.isTrain\n self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor\n self.load_dir = os.path.join(opt.checkpoints_dir, opt.load_dir)\n self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)\n\n def set_input(self, input):\n self.input = input\n\n def forward(self):\n pass\n\n # used in test time, no backprop\n def test(self):\n pass\n\n def get_image_paths(self):\n pass\n\n def optimize_parameters(self):\n pass\n\n def get_current_visuals(self):\n return self.input\n\n def get_current_errors(self):\n return {}\n\n def save(self, label):\n pass\n\n # helper saving function that can be used by subclasses\n def save_network(self, network, network_label, epoch_label, gpu_ids):\n save_filename = '%s_net_%s.pth' % (epoch_label, network_label)\n save_path = os.path.join(self.save_dir, save_filename)\n torch.save(network.cpu().state_dict(), save_path)\n if len(gpu_ids) and torch.cuda.is_available():\n network.cuda(gpu_ids[0])\n\n # helper loading function that can be used by subclasses\n def load_network(self, network, network_label, epoch_label):\n save_filename = '%s_net_%s.pth' % (epoch_label, network_label)\n save_path = os.path.join(self.load_dir, save_filename)\n network.load_state_dict(torch.load(save_path))\n\n # update learning rate (called once every epoch)\n def update_learning_rate(self):\n for scheduler in self.schedulers:\n scheduler.step()\n lr = self.optimizers[0].param_groups[0]['lr']\n print('learning rate = %.7f' % lr)\n\n def as_np(self, data):\n return data.cpu().data.numpy()\n" }, { "path": "drawing_and_style_transfer/models/networks.py", "content": "import torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport functools\nfrom torch.autograd import Variable\nfrom torch.optim import lr_scheduler\n\n\n###############################################################################\n# Functions\n###############################################################################\n\nclass pixel_norm(nn.Module):\n def forward(self, x, epsilon=1e-8):\n return x * torch.rsqrt(torch.mean(x.pow(2), dim=1, keepdim=True) + epsilon)\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n # print(classname)\n if classname.find('Conv') != -1:\n init.normal(m.weight.data, 0.0, 0.02)\n elif classname.find('Linear') != -1:\n init.normal(m.weight.data, 0.0, 0.02)\n elif classname.find('BatchNorm2d') != -1:\n init.normal(m.weight.data, 1.0, 0.02)\n init.constant(m.bias.data, 0.0)\n\n\ndef weights_init_xavier(m):\n classname = m.__class__.__name__\n # print(classname)\n if classname.find('Conv') != -1:\n init.xavier_normal(m.weight.data, gain=0.02)\n elif classname.find('Linear') != -1:\n init.xavier_normal(m.weight.data, gain=0.02)\n elif classname.find('BatchNorm2d') != -1:\n init.normal(m.weight.data, 1.0, 0.02)\n init.constant(m.bias.data, 0.0)\n\n\ndef weights_init_kaiming(m):\n classname = m.__class__.__name__\n # print(classname)\n if classname.find('Conv') != -1:\n init.kaiming_normal(m.weight.data, a=0, mode='fan_in')\n elif classname.find('Linear') != -1:\n init.kaiming_normal(m.weight.data, a=0, mode='fan_in')\n elif classname.find('BatchNorm2d') != -1:\n init.normal(m.weight.data, 1.0, 0.02)\n init.constant(m.bias.data, 0.0)\n\n\ndef weights_init_orthogonal(m):\n classname = m.__class__.__name__\n print(classname)\n if classname.find('Conv') != -1:\n init.orthogonal(m.weight.data, gain=1)\n elif classname.find('Linear') != -1:\n init.orthogonal(m.weight.data, gain=1)\n elif classname.find('BatchNorm2d') != -1:\n init.normal(m.weight.data, 1.0, 0.02)\n init.constant(m.bias.data, 0.0)\n\n\ndef init_weights(net, init_type='normal'):\n print('initialization method [%s]' % init_type)\n if init_type == 'normal':\n net.apply(weights_init_normal)\n elif init_type == 'xavier':\n net.apply(weights_init_xavier)\n elif init_type == 'kaiming':\n net.apply(weights_init_kaiming)\n elif init_type == 'orthogonal':\n net.apply(weights_init_orthogonal)\n else:\n raise NotImplementedError('initialization method [%s] is not implemented' % init_type)\n\n\ndef get_norm_layer(norm_type='instance'):\n if norm_type == 'batch':\n norm_layer = functools.partial(nn.BatchNorm2d, affine=True)\n elif norm_type == 'instance':\n norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)\n elif norm_type == 'none':\n norm_layer = None\n else:\n raise NotImplementedError('normalization layer [%s] is not found' % norm_type)\n return norm_layer\n\n\ndef get_scheduler(optimizer, opt):\n if opt.lr_policy == 'lambda':\n def lambda_rule(epoch):\n lr_l = 1.0 - max(0, epoch + 1 + opt.epoch_count - opt.niter) / float(opt.niter_decay + 1)\n return lr_l\n\n scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)\n elif opt.lr_policy == 'step':\n scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)\n elif opt.lr_policy == 'plateau':\n scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)\n else:\n return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)\n return scheduler\n\n\ndef define_ED(input_nc, output_nc, ngf, which_model_netG, norm='batch', use_dropout=False, init_type='normal',\n gpu_ids=[], n_downsampling=2, start=0, end=2, input_layer=True, output_layer=True, n_blocks_encoder=9,\n n_blocks_decoder=9, start_dec=0, end_dec=1):\n use_gpu = len(gpu_ids) > 0\n norm_layer = get_norm_layer(norm_type=norm)\n\n if use_gpu:\n assert (torch.cuda.is_available())\n\n if which_model_netG == 'resnet_9blocks':\n netE = ResnetEncoder(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout,\n n_blocks=n_blocks_encoder, gpu_ids=gpu_ids, n_downsampling=n_downsampling, start=start,\n end=end, input_layer=input_layer)\n netD = ResnetDecoder(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout,\n n_blocks=n_blocks_decoder, gpu_ids=gpu_ids, n_downsampling=n_downsampling, end=end_dec,\n start=start_dec, output_layer=output_layer)\n elif which_model_netG == 'resnet_6blocks':\n netE = ResnetEncoder(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout,\n n_blocks=n_blocks_encoder, gpu_ids=gpu_ids, n_downsampling=n_downsampling, start=start,\n end=end, input_layer=input_layer)\n netD = ResnetDecoder(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout,\n n_blocks=n_blocks_decoder, gpu_ids=gpu_ids, n_downsampling=n_downsampling, end=end_dec,\n start=start_dec, output_layer=output_layer)\n else:\n raise NotImplementedError('Generator model name [%s] is not recognized' % which_model_netG)\n if len(gpu_ids) > 0:\n netE.cuda(gpu_ids[0])\n netD.cuda(gpu_ids[0])\n init_weights(netE, init_type=init_type)\n init_weights(netD, init_type=init_type)\n return netE, netD\n\n\ndef define_G(input_nc, output_nc, ngf, which_model_netG, norm='batch', use_dropout=False, init_type='normal',\n gpu_ids=[]):\n netG = None\n use_gpu = len(gpu_ids) > 0\n norm_layer = get_norm_layer(norm_type=norm)\n\n if use_gpu:\n assert (torch.cuda.is_available())\n\n if which_model_netG == 'resnet_9blocks':\n netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9,\n gpu_ids=gpu_ids)\n elif which_model_netG == 'resnet_6blocks':\n netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6,\n gpu_ids=gpu_ids)\n elif which_model_netG == 'unet_128':\n netG = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout,\n gpu_ids=gpu_ids)\n elif which_model_netG == 'unet_256':\n netG = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout,\n gpu_ids=gpu_ids)\n else:\n raise NotImplementedError('Generator model name [%s] is not recognized' % which_model_netG)\n if len(gpu_ids) > 0:\n netG.cuda(gpu_ids[0])\n init_weights(netG, init_type=init_type)\n return netG\n\n\ndef define_D(input_nc, ndf, which_model_netD,\n n_layers_D=3, norm='batch', use_sigmoid=False, init_type='normal', gpu_ids=[]):\n use_gpu = len(gpu_ids) > 0\n norm_layer = get_norm_layer(norm_type=norm)\n\n if use_gpu:\n assert (torch.cuda.is_available())\n if which_model_netD == 'basic':\n netD = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer, use_sigmoid=use_sigmoid,\n gpu_ids=gpu_ids)\n elif which_model_netD == 'n_layers':\n netD = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer, use_sigmoid=use_sigmoid,\n gpu_ids=gpu_ids)\n elif which_model_netD == 'pixel':\n netD = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer, use_sigmoid=use_sigmoid, gpu_ids=gpu_ids)\n else:\n raise NotImplementedError('Discriminator model name [%s] is not recognized' %\n which_model_netD)\n if use_gpu:\n netD.cuda(gpu_ids[0])\n init_weights(netD, init_type=init_type)\n return netD\n\n\ndef print_network(net):\n num_params = 0\n for param in net.parameters():\n num_params += param.numel()\n print(net)\n print('Total number of parameters: %d' % num_params)\n\n\n##############################################################################\n# Classes\n##############################################################################\n\n\n# Defines the GAN loss which uses either LSGAN or the regular GAN.\n# When LSGAN is used, it is basically same as MSELoss,\n# but it abstracts away the need to create the target label tensor\n# that has the same size as the input\nclass GANLoss(nn.Module):\n def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0,\n tensor=torch.FloatTensor):\n super(GANLoss, self).__init__()\n self.real_label = target_real_label\n self.fake_label = target_fake_label\n self.real_label_var = None\n self.fake_label_var = None\n self.Tensor = tensor\n if use_lsgan:\n self.loss = nn.MSELoss()\n else:\n self.loss = nn.BCELoss()\n\n def get_target_tensor(self, input, target_is_real):\n target_tensor = None\n if target_is_real:\n create_label = ((self.real_label_var is None) or\n (self.real_label_var.numel() != input.numel()))\n if create_label:\n real_tensor = self.Tensor(input.size()).fill_(self.real_label)\n self.real_label_var = Variable(real_tensor, requires_grad=False)\n target_tensor = self.real_label_var\n else:\n create_label = ((self.fake_label_var is None) or\n (self.fake_label_var.numel() != input.numel()))\n if create_label:\n fake_tensor = self.Tensor(input.size()).fill_(self.fake_label)\n self.fake_label_var = Variable(fake_tensor, requires_grad=False)\n target_tensor = self.fake_label_var\n return target_tensor\n\n def __call__(self, input, target_is_real):\n target_tensor = self.get_target_tensor(input, target_is_real)\n return self.loss(input, target_tensor)\n\n # Defines the generator that consists of Resnet blocks between a few\n # downsampling/upsampling operations.\n # Code and idea originally from Justin Johnson's architecture.\n # https://github.com/jcjohnson/fast-neural-style/\n\n\nclass ResnetEncoder(nn.Module):\n def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False,\n gpu_ids=[], padding_type='reflect', n_downsampling=2, start=0, end=2, input_layer=True, n_blocks=6):\n assert (n_blocks >= 0)\n super(ResnetEncoder, self).__init__()\n self.input_nc = input_nc\n self.output_nc = output_nc\n self.ngf = ngf\n self.gpu_ids = gpu_ids\n if type(norm_layer) == functools.partial:\n use_bias = norm_layer.func == nn.InstanceNorm2d\n else:\n use_bias = norm_layer == nn.InstanceNorm2d\n\n model = []\n if input_layer:\n model = [nn.ReflectionPad2d(3),\n nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0,\n bias=use_bias),\n norm_layer(ngf),\n nn.ReLU(True)]\n\n for i in range(start, end):\n mult = 2 ** i\n model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3,\n stride=2, padding=1, bias=use_bias),\n norm_layer(ngf * mult * 2),\n nn.ReLU(True)]\n\n mult = 2 ** n_downsampling\n for i in range(n_blocks):\n model += [\n ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout,\n use_bias=use_bias)]\n\n self.model = nn.Sequential(*model)\n\n def forward(self, input):\n if self.gpu_ids and isinstance(input.data, torch.cuda.FloatTensor):\n return nn.parallel.data_parallel(self.model, input, self.gpu_ids)\n else:\n return self.model(input)\n\n\nclass ResnetDecoder(nn.Module):\n def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False,\n gpu_ids=[], padding_type='reflect', n_downsampling=2, start=0, end=2, output_layer=True, n_blocks=6):\n assert (n_blocks >= 0)\n super(ResnetDecoder, self).__init__()\n self.input_nc = input_nc\n self.output_nc = output_nc\n self.ngf = ngf\n self.gpu_ids = gpu_ids\n if type(norm_layer) == functools.partial:\n use_bias = norm_layer.func == nn.InstanceNorm2d\n else:\n use_bias = norm_layer == nn.InstanceNorm2d\n\n model = []\n mult = 2 ** n_downsampling\n for i in range(n_blocks):\n model += [\n ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout,\n use_bias=use_bias)]\n for i in range(start, end):\n mult = 2 ** (n_downsampling - i)\n model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),\n kernel_size=3, stride=2,\n padding=1, output_padding=1,\n bias=use_bias),\n norm_layer(int(ngf * mult / 2)),\n nn.ReLU(True)]\n\n if output_layer:\n model += [nn.ReflectionPad2d(3)]\n model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]\n model += [nn.Tanh()]\n\n self.model = nn.Sequential(*model)\n\n def forward(self, input):\n if self.gpu_ids and isinstance(input.data, torch.cuda.FloatTensor):\n return nn.parallel.data_parallel(self.model, input, self.gpu_ids)\n else:\n return self.model(input)\n\n\n# Defines the generator that consists of Resnet blocks between a few\n# downsampling/upsampling operations.\n# Code and idea originally from Justin Johnson's architecture.\n# https://github.com/jcjohnson/fast-neural-style/\nclass ResnetGenerator(nn.Module):\n def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6,\n gpu_ids=[], padding_type='reflect'):\n assert (n_blocks >= 0)\n super(ResnetGenerator, self).__init__()\n self.input_nc = input_nc\n self.output_nc = output_nc\n self.ngf = ngf\n self.gpu_ids = gpu_ids\n if type(norm_layer) == functools.partial:\n use_bias = norm_layer.func == nn.InstanceNorm2d\n else:\n use_bias = norm_layer == nn.InstanceNorm2d\n\n model = [nn.ReflectionPad2d(3),\n nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0,\n bias=use_bias),\n norm_layer(ngf),\n nn.ReLU(True)]\n\n n_downsampling = 2\n for i in range(n_downsampling):\n mult = 2 ** i\n model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3,\n stride=2, padding=1, bias=use_bias),\n norm_layer(ngf * mult * 2),\n nn.ReLU(True)]\n\n mult = 2 ** n_downsampling\n for i in range(n_blocks):\n model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout,\n use_bias=use_bias)]\n\n for i in range(n_downsampling):\n mult = 2 ** (n_downsampling - i)\n model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),\n kernel_size=3, stride=2,\n padding=1, output_padding=1,\n bias=use_bias),\n norm_layer(int(ngf * mult / 2)),\n nn.ReLU(True)]\n model += [nn.ReflectionPad2d(3)]\n model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]\n model += [nn.Tanh()]\n\n self.model = nn.Sequential(*model)\n\n def forward(self, input):\n if self.gpu_ids and isinstance(input.data, torch.cuda.FloatTensor):\n return nn.parallel.data_parallel(self.model, input, self.gpu_ids)\n else:\n return self.model(input)\n\n\n# Define a resnet block\nclass ResnetBlock(nn.Module):\n def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):\n super(ResnetBlock, self).__init__()\n self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)\n\n def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):\n conv_block = []\n p = 0\n if padding_type == 'reflect':\n conv_block += [nn.ReflectionPad2d(1)]\n elif padding_type == 'replicate':\n conv_block += [nn.ReplicationPad2d(1)]\n elif padding_type == 'zero':\n p = 1\n else:\n raise NotImplementedError('padding [%s] is not implemented' % padding_type)\n\n conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),\n norm_layer(dim),\n nn.ReLU(True)]\n if use_dropout:\n conv_block += [nn.Dropout(0.5)]\n\n p = 0\n if padding_type == 'reflect':\n conv_block += [nn.ReflectionPad2d(1)]\n elif padding_type == 'replicate':\n conv_block += [nn.ReplicationPad2d(1)]\n elif padding_type == 'zero':\n p = 1\n else:\n raise NotImplementedError('padding [%s] is not implemented' % padding_type)\n conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),\n norm_layer(dim)]\n\n return nn.Sequential(*conv_block)\n\n def forward(self, x):\n out = x + self.conv_block(x)\n return out\n\n\n# Defines the Unet generator.\n# |num_downs|: number of downsamplings in UNet. For example,\n# if |num_downs| == 7, image of size 128x128 will become of size 1x1\n# at the bottleneck\nclass UnetGenerator(nn.Module):\n def __init__(self, input_nc, output_nc, num_downs, ngf=64,\n norm_layer=nn.BatchNorm2d, use_dropout=False, gpu_ids=[]):\n super(UnetGenerator, self).__init__()\n self.gpu_ids = gpu_ids\n\n # construct unet structure\n unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer,\n innermost=True)\n for i in range(num_downs - 5):\n unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block,\n norm_layer=norm_layer, use_dropout=use_dropout)\n unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block,\n norm_layer=norm_layer)\n unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block,\n norm_layer=norm_layer)\n unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)\n unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True,\n norm_layer=norm_layer)\n\n self.model = unet_block\n\n def forward(self, input):\n if self.gpu_ids and isinstance(input.data, torch.cuda.FloatTensor):\n return nn.parallel.data_parallel(self.model, input, self.gpu_ids)\n else:\n return self.model(input)\n\n\n# Defines the submodule with skip connection.\n# X -------------------identity---------------------- X\n# |-- downsampling -- |submodule| -- upsampling --|\nclass UnetSkipConnectionBlock(nn.Module):\n def __init__(self, outer_nc, inner_nc, input_nc=None,\n submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):\n super(UnetSkipConnectionBlock, self).__init__()\n self.outermost = outermost\n if type(norm_layer) == functools.partial:\n use_bias = norm_layer.func == nn.InstanceNorm2d\n else:\n use_bias = norm_layer == nn.InstanceNorm2d\n if input_nc is None:\n input_nc = outer_nc\n downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,\n stride=2, padding=1, bias=use_bias)\n downrelu = nn.LeakyReLU(0.2, True)\n downnorm = norm_layer(inner_nc)\n uprelu = nn.ReLU(True)\n upnorm = norm_layer(outer_nc)\n\n if outermost:\n upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,\n kernel_size=4, stride=2,\n padding=1)\n down = [downconv]\n up = [uprelu, upconv, nn.Tanh()]\n model = down + [submodule] + up\n elif innermost:\n upconv = nn.ConvTranspose2d(inner_nc, outer_nc,\n kernel_size=4, stride=2,\n padding=1, bias=use_bias)\n down = [downrelu, downconv]\n up = [uprelu, upconv, upnorm]\n model = down + up\n else:\n upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,\n kernel_size=4, stride=2,\n padding=1, bias=use_bias)\n down = [downrelu, downconv, downnorm]\n up = [uprelu, upconv, upnorm]\n\n if use_dropout:\n model = down + [submodule] + up + [nn.Dropout(0.5)]\n else:\n model = down + [submodule] + up\n\n self.model = nn.Sequential(*model)\n\n def forward(self, x):\n if self.outermost:\n return self.model(x)\n else:\n return torch.cat([x, self.model(x)], 1)\n\n\n# Defines the PatchGAN discriminator with the specified arguments.\nclass NLayerDiscriminator(nn.Module):\n def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False, gpu_ids=[]):\n super(NLayerDiscriminator, self).__init__()\n self.gpu_ids = gpu_ids\n if type(norm_layer) == functools.partial:\n use_bias = norm_layer.func == nn.InstanceNorm2d\n else:\n use_bias = norm_layer == nn.InstanceNorm2d\n\n kw = 4\n padw = 1\n sequence = [\n nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),\n nn.LeakyReLU(0.2, True)\n ]\n\n nf_mult = 1\n nf_mult_prev = 1\n for n in range(1, n_layers):\n nf_mult_prev = nf_mult\n nf_mult = min(2 ** n, 8)\n sequence += [\n nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,\n kernel_size=kw, stride=2, padding=padw, bias=use_bias),\n norm_layer(ndf * nf_mult),\n nn.LeakyReLU(0.2, True)\n ]\n\n nf_mult_prev = nf_mult\n nf_mult = min(2 ** n_layers, 8)\n sequence += [\n nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,\n kernel_size=kw, stride=1, padding=padw, bias=use_bias),\n norm_layer(ndf * nf_mult),\n nn.LeakyReLU(0.2, True)\n ]\n\n sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]\n\n if use_sigmoid:\n sequence += [nn.Sigmoid()]\n\n self.model = nn.Sequential(*sequence)\n\n def forward(self, input):\n if len(self.gpu_ids) and isinstance(input.data, torch.cuda.FloatTensor):\n return nn.parallel.data_parallel(self.model, input, self.gpu_ids)\n else:\n return self.model(input)\n\n\nclass PixelDiscriminator(nn.Module):\n def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d, use_sigmoid=False, gpu_ids=[]):\n super(PixelDiscriminator, self).__init__()\n self.gpu_ids = gpu_ids\n if type(norm_layer) == functools.partial:\n use_bias = norm_layer.func == nn.InstanceNorm2d\n else:\n use_bias = norm_layer == nn.InstanceNorm2d\n\n self.net = [\n nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),\n nn.LeakyReLU(0.2, True),\n nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),\n norm_layer(ndf * 2),\n nn.LeakyReLU(0.2, True),\n nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]\n\n if use_sigmoid:\n self.net.append(nn.Sigmoid())\n\n self.net = nn.Sequential(*self.net)\n\n def forward(self, input):\n if len(self.gpu_ids) and isinstance(input.data, torch.cuda.FloatTensor):\n return nn.parallel.data_parallel(self.net, input, self.gpu_ids)\n else:\n return self.net(input)\n" }, { "path": "drawing_and_style_transfer/models/ost.py", "content": "import torch\nfrom collections import OrderedDict\nfrom torch.autograd import Variable\nimport itertools\nimport util.util as util\nfrom util.image_pool import ImagePool\nfrom .base_model import BaseModel\nfrom . import networks\n\n\nclass OSTModel(BaseModel):\n def name(self):\n return 'OSTModel'\n\n def _compute_kl(self, mu):\n mu_2 = torch.pow(mu, 2)\n encoding_loss = torch.mean(mu_2)\n return encoding_loss\n\n def set_encoders_and_decoders(self, opt):\n n_downsampling = opt.n_downsampling\n start_unshared = 0\n num_unshared = opt.num_unshared\n start_shared = num_unshared\n end_shared = n_downsampling\n start_dec_shared = start_unshared\n end_dec_shared = start_unshared + (end_shared - start_shared)\n start_dec_unshared = end_dec_shared\n end_dec_unshared = n_downsampling\n\n num_res_blocks_unshared = opt.num_res_blocks_unshared\n n_res_blocks_shared = opt.num_res_blocks_shared\n\n self.netEnc_a, self.netDec_a = networks.define_ED(opt.input_nc, opt.output_nc,\n opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout,\n opt.init_type, self.gpu_ids,\n n_blocks_encoder=num_res_blocks_unshared,\n n_blocks_decoder=num_res_blocks_unshared,\n start=start_unshared,\n end=num_unshared, n_downsampling=n_downsampling,\n input_layer=True, output_layer=True,\n start_dec=start_dec_unshared, end_dec=end_dec_unshared)\n\n self.netEnc_b, self.netDec_b = networks.define_ED(opt.input_nc, opt.output_nc,\n opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout,\n opt.init_type, self.gpu_ids,\n n_blocks_encoder=num_res_blocks_unshared,\n n_blocks_decoder=num_res_blocks_unshared,\n start=start_unshared,\n end=num_unshared, n_downsampling=n_downsampling,\n input_layer=True,\n output_layer=True, start_dec=start_dec_unshared,\n end_dec=end_dec_unshared)\n\n self.netEnc_shared, self.netDec_shared = networks.define_ED(opt.input_nc, opt.output_nc,\n opt.ngf, opt.which_model_netG, opt.norm,\n not opt.no_dropout,\n opt.init_type, self.gpu_ids,\n n_blocks_encoder=n_res_blocks_shared,\n n_blocks_decoder=n_res_blocks_shared,\n start=start_shared, n_downsampling=n_downsampling,\n end=end_shared,\n input_layer=False,\n output_layer=False, start_dec=start_dec_shared,\n end_dec=end_dec_shared)\n\n def initialize(self, opt):\n BaseModel.initialize(self, opt)\n self.set_encoders_and_decoders(opt)\n\n if self.isTrain:\n use_sigmoid = opt.no_lsgan\n self.netD_a = networks.define_D(opt.output_nc, opt.ndf,\n opt.which_model_netD,\n opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, self.gpu_ids)\n self.netD_b = networks.define_D(opt.output_nc, opt.ndf,\n opt.which_model_netD,\n opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, self.gpu_ids)\n\n if not opt.dont_load_pretrained_autoencoder:\n which_epoch = opt.which_epoch\n self.load_network(self.netEnc_b, 'Enc_b', which_epoch)\n self.load_network(self.netDec_b, 'Dec_b', which_epoch)\n self.load_network(self.netEnc_shared, 'Enc_shared', which_epoch)\n self.load_network(self.netDec_shared, 'Dec_shared', which_epoch)\n\n if not self.isTrain or opt.continue_train:\n which_epoch = opt.which_epoch\n self.load_network(self.netEnc_a, 'Enc_a', which_epoch)\n self.load_network(self.netDec_a, 'Dec_a', which_epoch)\n self.load_network(self.netEnc_b, 'Enc_b', which_epoch)\n self.load_network(self.netDec_b, 'Dec_b', which_epoch)\n self.load_network(self.netEnc_shared, 'Enc_shared', which_epoch)\n self.load_network(self.netDec_shared, 'Dec_shared', which_epoch)\n if self.isTrain:\n self.load_network(self.netD_a, 'D_a', which_epoch)\n self.load_network(self.netD_b, 'D_b', which_epoch)\n\n if self.isTrain:\n self.fake_A_pool = ImagePool(opt.pool_size)\n self.fake_B_pool = ImagePool(opt.pool_size)\n\n # define loss functions\n self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)\n self.criterionIdt = torch.nn.L1Loss()\n\n # initialize optimizers\n self.optimizer_Enc_a = torch.optim.Adam(self.netEnc_a.parameters(),\n lr=opt.lr, betas=(opt.beta1, 0.999))\n self.optimizer_Dec_a = torch.optim.Adam(self.netDec_a.parameters(),\n lr=opt.lr, betas=(opt.beta1, 0.999))\n self.optimizer_Enc_b = torch.optim.Adam(\n itertools.chain(self.netEnc_b.parameters(), self.netEnc_shared.parameters()),\n lr=opt.lr, betas=(opt.beta1, 0.999))\n self.optimizer_Dec_b = torch.optim.Adam(\n itertools.chain(self.netDec_b.parameters(), self.netDec_shared.parameters()),\n lr=opt.lr, betas=(opt.beta1, 0.999))\n self.optimizer_D_a = torch.optim.Adam(self.netD_a.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))\n self.optimizer_D_b = torch.optim.Adam(self.netD_b.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))\n\n self.optimizers = []\n self.schedulers = []\n self.optimizers.append(self.optimizer_Enc_a)\n self.optimizers.append(self.optimizer_Dec_a)\n self.optimizers.append(self.optimizer_Enc_b)\n self.optimizers.append(self.optimizer_Dec_b)\n self.optimizers.append(self.optimizer_D_a)\n self.optimizers.append(self.optimizer_D_b)\n for optimizer in self.optimizers:\n self.schedulers.append(networks.get_scheduler(optimizer, opt))\n\n print('---------- Networks initialized -------------')\n networks.print_network(self.netEnc_a)\n networks.print_network(self.netDec_a)\n networks.print_network(self.netEnc_b)\n networks.print_network(self.netDec_b)\n networks.print_network(self.netEnc_shared)\n networks.print_network(self.netDec_shared)\n if self.isTrain:\n networks.print_network(self.netD_a)\n networks.print_network(self.netD_b)\n print('-----------------------------------------------')\n\n def set_input(self, input):\n AtoB = self.opt.which_direction == 'AtoB'\n input_A = input['A' if AtoB else 'B']\n input_B = input['B' if AtoB else 'A']\n if len(self.gpu_ids) > 0:\n input_A = input_A.cuda(self.gpu_ids[0], async=True)\n input_B = input_B.cuda(self.gpu_ids[0], async=True)\n self.input_A = input_A\n self.input_B = input_B\n self.image_paths = input['A_paths' if AtoB else 'B_paths']\n\n def forward(self):\n self.real_A = Variable(self.input_A)\n self.real_B = Variable(self.input_B)\n\n def test(self):\n real_A = Variable(self.input_A, volatile=True)\n real_B = Variable(self.input_B, volatile=True)\n enc_a = self.netEnc_shared(self.netEnc_a(real_A))\n enc_b = self.netEnc_shared(self.netEnc_b(real_B))\n\n fake_AA = self.netDec_a(self.netDec_shared(enc_a))\n fake_AB = self.netDec_b(self.netDec_shared(enc_a))\n fake_BB = self.netDec_b(self.netDec_shared(enc_b))\n\n enc_ab = self.netEnc_shared(self.netEnc_b(fake_AB))\n fake_ABA = self.netDec_a(self.netDec_shared(enc_ab))\n\n self.fake_AA = fake_AA.data\n self.fake_AB = fake_AB.data\n self.fake_BB = fake_BB.data\n self.fake_ABA = fake_ABA.data\n\n # get image paths\n def get_image_paths(self):\n return self.image_paths\n\n def backward_D_basic(self, netD, real, fake):\n # Real\n pred_real = netD(real)\n loss_D_real = self.criterionGAN(pred_real, True)\n # Fake\n pred_fake = netD(fake.detach())\n loss_D_fake = self.criterionGAN(pred_fake, False)\n # Combined loss\n loss_D = (loss_D_real + loss_D_fake) * 0.5\n # backward\n loss_D.backward()\n return loss_D\n\n def backward_D(self):\n fake_AB = self.fake_B_pool.query(self.fake_AB)\n loss_D_ab = self.backward_D_basic(self.netD_b, self.real_B, fake_AB)\n self.loss_D_ab = loss_D_ab.data[0]\n\n fake_BB = self.fake_B_pool.query(self.fake_BB)\n loss_D_bb = self.backward_D_basic(self.netD_b, self.real_B, fake_BB)\n self.loss_D_bb = loss_D_bb.data[0]\n\n def backward_G(self):\n # GAN loss D_A(G_A(A))\n enc_a = self.netEnc_shared(self.netEnc_a(self.real_A))\n enc_b = self.netEnc_shared(self.netEnc_b(self.real_B))\n fake_AA = self.netDec_a(self.netDec_shared(enc_a))\n fake_AB = self.netDec_b(self.netDec_shared(enc_a))\n fake_BB = self.netDec_b(self.netDec_shared(enc_b))\n enc_ab = self.netEnc_shared(self.netEnc_b(fake_AB))\n fake_ABA = self.netDec_a(self.netDec_shared(enc_ab))\n\n pred_fake_AB = self.netD_b(fake_AB)\n loss_Gan_AB = self.criterionGAN(pred_fake_AB, True)\n loss_idt_A = self.criterionIdt(fake_AA, self.real_A)\n loss_cycle_A = self.opt.lambda_A * self.criterionIdt(fake_ABA, self.real_A)\n loss_idt_B = self.criterionIdt(fake_BB, self.real_B)\n pred_fake_BB = self.netD_b(fake_BB)\n loss_Gan_BB = self.criterionGAN(pred_fake_BB, True)\n loss_kl_B = self.opt.kl_lambda * self._compute_kl(enc_b)\n\n # combined losses\n loss_G_B = loss_idt_B + loss_kl_B + loss_Gan_BB\n loss_G_A = loss_Gan_AB + loss_cycle_A + loss_idt_A\n\n self.fake_AA = fake_AA.data\n self.fake_BB = fake_BB.data\n self.fake_AB = fake_AB.data\n self.fake_ABA = fake_ABA.data\n self.loss_cycle_A = loss_cycle_A.data[0]\n self.loss_Gan_AB = loss_Gan_AB.data[0]\n self.loss_Gan_BB = loss_Gan_AB.data[0]\n self.loss_idt_A = loss_idt_A.data[0]\n self.loss_idt_B = loss_idt_B.data[0]\n self.loss_kl_B = loss_kl_B.data[0]\n\n return loss_G_A, loss_G_B\n\n def optimize_parameters(self):\n # forward\n self.forward()\n loss_G_A, loss_G_B = self.backward_G()\n\n # x loss updates\n self.optimizer_Enc_a.zero_grad()\n self.optimizer_Dec_a.zero_grad()\n loss_G_A.backward(retain_graph=True)\n self.optimizer_Enc_a.step()\n self.optimizer_Dec_a.step()\n\n # B loss updates\n self.optimizer_Enc_b.zero_grad()\n self.optimizer_Dec_b.zero_grad()\n loss_G_B.backward()\n self.optimizer_Enc_b.step()\n self.optimizer_Dec_b.step()\n\n # D\n self.optimizer_D_a.zero_grad()\n self.optimizer_D_b.zero_grad()\n self.backward_D()\n self.optimizer_D_a.step()\n self.optimizer_D_b.step()\n\n def get_current_errors(self):\n ret_errors = OrderedDict(\n [('D_ab', self.loss_D_ab), ('D_bb', self.loss_D_bb),\n ('G_AB', self.loss_Gan_AB), ('G_BB', self.loss_Gan_BB),\n ('Idt_B', self.loss_idt_B), ('Idt_A', self.loss_idt_A),\n ('Cycle_A', self.loss_cycle_A), ('Kl_B', self.loss_kl_B), ])\n return ret_errors\n\n def get_current_visuals(self):\n real_A = util.tensor2im(self.input_A)\n real_B = util.tensor2im(self.input_B)\n fake_BB = util.tensor2im(self.fake_BB)\n fake_AB = util.tensor2im(self.fake_AB)\n fake_AA = util.tensor2im(self.fake_AA)\n fake_ABA = util.tensor2im(self.fake_ABA)\n\n ret_visuals = OrderedDict(\n [('real_B', real_B), ('fake_BB', fake_BB),\n ('real_A', real_A), ('fake_AA', fake_AA), ('fake_AB', fake_AB), ('fake_ABA', fake_ABA), ])\n return ret_visuals\n\n def save(self, label):\n self.save_network(self.netEnc_a, 'Enc_a', label, self.gpu_ids)\n self.save_network(self.netDec_a, 'Dec_a', label, self.gpu_ids)\n self.save_network(self.netD_a, 'D_a', label, self.gpu_ids)\n self.save_network(self.netEnc_b, 'Enc_b', label, self.gpu_ids)\n self.save_network(self.netDec_b, 'Dec_b', label, self.gpu_ids)\n self.save_network(self.netD_b, 'D_b', label, self.gpu_ids)\n self.save_network(self.netEnc_shared, 'Enc_shared', label, self.gpu_ids)\n self.save_network(self.netDec_shared, 'Dec_shared', label, self.gpu_ids)\n" }, { "path": "drawing_and_style_transfer/models/test_model.py", "content": "from torch.autograd import Variable\nfrom collections import OrderedDict\nimport util.util as util\nfrom .base_model import BaseModel\nfrom . import networks\n\n\nclass TestModel(BaseModel):\n def name(self):\n return 'TestModel'\n\n def initialize(self, opt):\n assert (not opt.isTrain)\n BaseModel.initialize(self, opt)\n self.netG = networks.define_G(opt.input_nc, opt.output_nc,\n opt.ngf, opt.which_model_netG,\n opt.norm, not opt.no_dropout,\n opt.init_type,\n self.gpu_ids)\n which_epoch = opt.which_epoch\n self.load_network(self.netG, 'G', which_epoch)\n\n print('---------- Networks initialized -------------')\n networks.print_network(self.netG)\n print('-----------------------------------------------')\n\n def set_input(self, input):\n # we need to use single_dataset mode\n input_A = input['A']\n if len(self.gpu_ids) > 0:\n input_A = input_A.cuda(self.gpu_ids[0], async=True)\n self.input_A = input_A\n self.image_paths = input['A_paths']\n\n def test(self):\n self.real_A = Variable(self.input_A, volatile=True)\n self.fake_B = self.netG(self.real_A)\n\n # get image paths\n def get_image_paths(self):\n return self.image_paths\n\n def get_current_visuals(self):\n real_A = util.tensor2im(self.real_A.data)\n fake_B = util.tensor2im(self.fake_B.data)\n return OrderedDict([('real_A', real_A), ('fake_B', fake_B)])\n" }, { "path": "drawing_and_style_transfer/options/__init__.py", "content": "" }, { "path": "drawing_and_style_transfer/options/base_options.py", "content": "import argparse\nimport os\nfrom util import util\nimport torch\n\n\nclass BaseOptions():\n def __init__(self):\n self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n self.initialized = False\n\n def initialize(self):\n self.parser.add_argument('--dataroot', required=True,\n help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')\n self.parser.add_argument('--batchSize', type=int, default=1, help='input batch size')\n self.parser.add_argument('--max_items_A', type=int, default=-1,\n help='max number of items for domain A, -1 indicates no maximum')\n self.parser.add_argument('--max_items_B', type=int, default=-1,\n help='max number of items for domain B, -1 indicates no maximum')\n self.parser.add_argument('--start', type=int, default=0,\n help='starting index of items of domain A, after sorting')\n self.parser.add_argument('--loadSize', type=int, default=286, help='scale images to this size')\n self.parser.add_argument('--fineSize', type=int, default=256, help='then crop to this size')\n self.parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels')\n self.parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels')\n self.parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')\n self.parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer')\n self.parser.add_argument('--which_model_netD', type=str, default='basic', help='selects model to use for netD')\n self.parser.add_argument('--which_model_netG', type=str, default='resnet_9blocks',\n help='selects model to use for netG and netED')\n self.parser.add_argument('--n_layers_D', type=int, default=3, help='only used if which_model_netD==n_layers')\n self.parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU')\n self.parser.add_argument('--name', type=str, default='experiment_name',\n help='name of the experiment. It decides where to store samples and models')\n self.parser.add_argument('--dataset_mode', type=str, default='unaligned',\n help='chooses how datasets are loaded. [unaligned | aligned | single]')\n self.parser.add_argument('--model', type=str, default='ost',\n help='chooses which model to use. ost, autoencoder, test.')\n self.parser.add_argument('--which_direction', type=str, default='AtoB', help='AtoB or BtoA')\n self.parser.add_argument('--nThreads', default=2, type=int, help='# threads for loading data')\n self.parser.add_argument('--load_dir', type=str, default='./checkpoints', help='models are loaded here')\n self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')\n self.parser.add_argument('--norm', type=str, default='instance',\n help='instance normalization or batch normalization')\n self.parser.add_argument('--serial_batches', action='store_true',\n help='if true, takes images in order to make batches, otherwise takes them randomly')\n self.parser.add_argument('--display_winsize', type=int, default=256, help='display window size')\n self.parser.add_argument('--display_id', type=int, default=1, help='window id of the web display')\n self.parser.add_argument('--display_server', type=str, default=\"http://localhost\",\n help='visdom server of the web display')\n self.parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display')\n self.parser.add_argument('--no_dropout', action='store_true', help='no dropout for the generator')\n self.parser.add_argument('--resize_or_crop', type=str, default='resize_and_crop',\n help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop]')\n self.parser.add_argument('--no_flip_and_rotation', action='store_true',\n help='if specified, do not flip and rotate the images for data augmentation')\n self.parser.add_argument('--rotation_degree', type=int, default=7, help='rotation degree used for augmentation')\n self.parser.add_argument('--init_type', type=str, default='normal',\n help='network initialization [normal|xavier|kaiming|orthogonal]')\n self.parser.add_argument('--A', type=str, default='A',\n help='used to exchange dataset A for B by setting the value to B')\n self.parser.add_argument('--B', type=str, default='B',\n help='used to exchange dataset B for A by setting the value to A')\n self.parser.add_argument('--n_downsampling', type=int, default=2,\n help=\"number of downsampling/upsampling convolutional/deconvolutional layers\")\n self.parser.add_argument('--num_unshared', type=int, default=1,\n help=\"number of unshared encoder/decoder layers, not including input and final layers\")\n self.parser.add_argument('--num_res_blocks_unshared', type=int, default=0,\n help='number of unshared resnet blocks')\n self.parser.add_argument('--num_res_blocks_shared', type=int, default=6, help='number of shared resnet blocks')\n\n self.initialized = True\n\n def parse(self):\n if not self.initialized:\n self.initialize()\n self.opt = self.parser.parse_args()\n self.opt.isTrain = self.isTrain # train or test\n\n str_ids = self.opt.gpu_ids.split(',')\n self.opt.gpu_ids = []\n for str_id in str_ids:\n id = int(str_id)\n if id >= 0:\n self.opt.gpu_ids.append(id)\n\n # set gpu ids\n if len(self.opt.gpu_ids) > 0:\n torch.cuda.set_device(self.opt.gpu_ids[0])\n\n args = vars(self.opt)\n\n print('------------ Options -------------')\n for k, v in sorted(args.items()):\n print('%s: %s' % (str(k), str(v)))\n print('-------------- End ----------------')\n\n # save to the disk\n expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name)\n util.mkdirs(expr_dir)\n file_name = os.path.join(expr_dir, 'opt.txt')\n with open(file_name, 'wt') as opt_file:\n opt_file.write('------------ Options -------------\\n')\n for k, v in sorted(args.items()):\n opt_file.write('%s: %s\\n' % (str(k), str(v)))\n opt_file.write('-------------- End ----------------\\n')\n return self.opt\n" }, { "path": "drawing_and_style_transfer/options/test_options.py", "content": "from .base_options import BaseOptions\n\n\nclass TestOptions(BaseOptions):\n def initialize(self):\n BaseOptions.initialize(self)\n self.parser.add_argument('--ntest', type=int, default=float(\"inf\"), help='# of test examples.')\n self.parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')\n self.parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')\n self.parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')\n self.parser.add_argument('--which_epoch', type=str, default='latest',\n help='which epoch to load? set to latest to use latest cached model')\n self.parser.add_argument('--how_many', type=int, default=50, help='how many test images to run')\n self.isTrain = False\n" }, { "path": "drawing_and_style_transfer/options/train_options.py", "content": "from .base_options import BaseOptions\n\n\nclass TrainOptions(BaseOptions):\n def initialize(self):\n BaseOptions.initialize(self)\n self.parser.add_argument('--display_freq', type=int, default=100,\n help='frequency of showing training results on screen')\n self.parser.add_argument('--display_single_pane_ncols', type=int, default=0,\n help='if positive, display all images in a single visdom web panel with certain number of images per row.')\n self.parser.add_argument('--update_html_freq', type=int, default=1000,\n help='frequency of saving training results to html')\n self.parser.add_argument('--print_freq', type=int, default=100,\n help='frequency of showing training results on console')\n self.parser.add_argument('--save_latest_freq', type=int, default=10000,\n help='frequency of saving the latest results')\n self.parser.add_argument('--save_epoch_freq', type=int, default=10,\n help='frequency of saving checkpoints at the end of epochs')\n self.parser.add_argument('--continue_train', action='store_true',\n help='continue training: load the latest model')\n self.parser.add_argument('--epoch_count', type=int, default=1,\n help='the starting epoch count, we save the model by , +, ...')\n self.parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')\n self.parser.add_argument('--which_epoch', type=str, default='latest',\n help='which epoch to load? set to latest to use latest cached model')\n self.parser.add_argument('--niter', type=int, default=60, help='# of iter at starting learning rate')\n self.parser.add_argument('--niter_decay', type=int, default=20,\n help='# of iter to linearly decay learning rate to zero')\n self.parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')\n self.parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')\n self.parser.add_argument('--no_lsgan', action='store_true',\n help='do *not* use least square GAN, if false, use vanilla GAN')\n self.parser.add_argument('--kl_lambda', type=float, default=0.1, help='weight for kl loss')\n self.parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)')\n self.parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')\n self.parser.add_argument('--pool_size', type=int, default=50,\n help='the size of image buffer that stores previously generated images')\n self.parser.add_argument('--dont_load_pretrained_autoencoder', action='store_true',\n help='do not load pretrained autoencoder')\n self.parser.add_argument('--no_html', action='store_true',\n help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/')\n self.parser.add_argument('--lr_policy', type=str, default='lambda',\n help='learning rate policy: lambda|step|plateau')\n self.parser.add_argument('--lr_decay_iters', type=int, default=50,\n help='multiply by a gamma every lr_decay_iters iterations')\n\n self.isTrain = True\n" }, { "path": "drawing_and_style_transfer/scripts/test_ost.sh", "content": "# images to cityscapes\npython test.py --dataroot=./datasets/cityscapes/ --name=cityscapes_ost --load_dir=cityscapes_ost --model=ost --no_dropout --n_downsampling=3 --num_unshared=3 --start=0 --max_items_A=1\n# cityscapes to images\npython test.py --dataroot=./datasets/cityscapes/ --name=cityscapes_ost_reverse --load_dir=cityscapes_ost_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1 --A='B' --B='A'\n\n# images to facades\npython test.py --dataroot=./datasets/facades/ --name=facades_ost --load_dir=facades_ost --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1\n# facades to images\npython test.py --dataroot=./datasets/facades/ --name=facades_ost_reverse --load_dir=facades_ost_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1 --A='B' --B='A'\n\n# aerial view to maps\npython test.py --dataroot=./datasets/maps/ --name=maps_ost --load_dir=maps_ost --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1\n# maps to aerial view\npython test.py --dataroot=./datasets/maps/ --name=maps_ost_reverse --load_dir=maps_ost_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1 --A='B' --B='A'\n\n# monet2photo\npython test.py --dataroot=./datasets/monet2photo/ --name=monet2photo_ost --load_dir=monet2photo_ost --model=ost --no_dropout --n_downsampling=2 --num_unshared=0 --start=0 --max_items_A=1\n# photo2monet\npython test.py --dataroot=./datasets/monet2photo/ --name=monet2photo_ost_reverse --load_dir=monet2photo_ost_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=0 --start=0 --max_items_A=1 --A='B' --B='A'\n\n# summer2winter\npython test.py --dataroot=./datasets/summer2winter_yosemite/ --name=summer2winter_yosemite_ost --load_dir=summer2winter_yosemite_ost --model=ost --no_dropout --n_downsampling=2 --num_unshared=0 --start=0 --max_items_A=1\n# winter2summer\npython test.py --dataroot=./datasets/summer2winter_yosemite/ --name=summer2winter_yosemite_ost_reverse --load_dir=summer2winter_yosemite_ost_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=0 --start=0 --max_items_A=1 --A='B' --B='A'" }, { "path": "drawing_and_style_transfer/scripts/train_autoencoder.sh", "content": "# images to cityscapes\npython train.py --dataroot=./datasets/cityscapes/trainB --name=cityscapes_autoencoder --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=3 --num_unshared=3\n# cityscapes to images\npython train.py --dataroot=./datasets/cityscapes/trainA --name=cityscapes_autoencoder_reverse --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=2\n\n# images to facades\npython train.py --dataroot=./datasets/facades/trainB --name=facades_autoencoder --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=2\n# facades to images\npython train.py --dataroot=./datasets/facades/trainA --name=facades_autoencoder_reverse --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=2\n\n# aerial view to maps\npython train.py --dataroot=./datasets/maps/trainB --name=maps_autoencoder --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=2\n# maps to aerial view\npython train.py --dataroot=./datasets/maps/trainA --name=maps_autoencoder_reverse --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=2\n\n# monet2photo\npython train.py --dataroot=./datasets/monet2photo/trainB --name=monet2photo_autoencoder --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=0\n# photo2monet\npython train.py --dataroot=./datasets/monet2photo/trainA --name=monet2photo_autoencoder_reverse --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=0\n\n# summer2winter\npython train.py --dataroot=./datasets/summer2winter_yosemite/trainB --name=summer2winter_yosemite_autoencoder --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=0\n# winter2summer\npython train.py --dataroot=./datasets/summer2winter_yosemite/trainA --name=summer2winter_yosemite_autoencoder_reverse --model=autoencoder --dataset_mode=single --no_dropout --n_downsampling=2 --num_unshared=0" }, { "path": "drawing_and_style_transfer/scripts/train_ost.sh", "content": "# images to cityscapes\npython train.py --dataroot=./datasets/cityscapes/ --name=cityscapes_ost --load_dir=cityscapes_autoencoder --model=ost --no_dropout --n_downsampling=3 --num_unshared=3 --start=0 --max_items_A=1\n# cityscapes to images\npython train.py --dataroot=./datasets/cityscapes/ --name=cityscapes_ost_reverse --load_dir=cityscapes_autoencoder_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1 --A='B' --B='A'\n\n# images to facades\npython train.py --dataroot=./datasets/facades/ --name=facades_ost --load_dir=facades_autoencoder --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1\n# facades to images\npython train.py --dataroot=./datasets/facades/ --name=facades_ost_reverse --load_dir=facades_autoencoder_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1 --A='B' --B='A'\n\n# aerial view to maps\npython train.py --dataroot=./datasets/maps/ --name=maps_ost --load_dir=maps_autoencoder --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1\n# maps to aerial view\npython train.py --dataroot=./datasets/maps/ --name=maps_ost_reverse --load_dir=maps_autoencoder_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=2 --start=0 --max_items_A=1 --A='B' --B='A'\n\n# monet2photo\npython train.py --dataroot=./datasets/monet2photo/ --name=monet2photo_ost --load_dir=monet2photo_autoencoder --model=ost --no_dropout --n_downsampling=2 --num_unshared=0 --start=0 --max_items_A=1\n# photo2monet\npython train.py --dataroot=./datasets/monet2photo/ --name=monet2photo_ost_reverse --load_dir=monet2photo_autoencoder_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=0 --start=0 --max_items_A=1 --A='B' --B='A'\n\n# summer2winter\npython train.py --dataroot=./datasets/summer2winter_yosemite/ --name=summer2winter_yosemite_ost --load_dir=summer2winter_yosemite_autoencoder --model=ost --no_dropout --n_downsampling=2 --num_unshared=0 --start=0 --max_items_A=1\n# winter2summer\npython train.py --dataroot=./datasets/summer2winter_yosemite/ --name=summer2winter_yosemite_ost_reverse --load_dir=summer2winter_yosemite_autoencoder_reverse --model=ost --no_dropout --n_downsampling=2 --num_unshared=0 --start=0 --max_items_A=1 --A='B' --B='A'" }, { "path": "drawing_and_style_transfer/test.py", "content": "import os\nfrom options.test_options import TestOptions\nfrom data import CreateDataLoader\nfrom models import create_model\nfrom util.visualizer import Visualizer\nfrom util import html\n\nif __name__ == '__main__':\n opt = TestOptions().parse()\n opt.nThreads = 1 # test code only supports nThreads = 1\n opt.batchSize = 1 # test code only supports batchSize = 1\n opt.serial_batches = True # no shuffle\n opt.no_flip = True # no flip\n\n # We are interested on testing only on x's in A on which we trained\n opt.phase = 'train'\n if opt.max_items_A >= 0:\n opt.max_items_B = opt.max_items_A\n if opt.max_items_B >= 0:\n opt.max_items_A = opt.max_items_B\n\n data_loader = CreateDataLoader(opt)\n dataset = data_loader.load_data()\n model = create_model(opt)\n visualizer = Visualizer(opt)\n # create website\n web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))\n webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))\n # test\n for i, data in enumerate(dataset):\n if i >= opt.how_many:\n break\n model.set_input(data)\n model.test()\n visuals = model.get_current_visuals()\n img_path = model.get_image_paths()\n print('%04d: process image... %s' % (i, img_path))\n visualizer.save_images(webpage, visuals, img_path, aspect_ratio=opt.aspect_ratio, index=i, split=0)\n\n webpage.save()\n" }, { "path": "drawing_and_style_transfer/train.py", "content": "import time\nfrom options.train_options import TrainOptions\nfrom data import CreateDataLoader\nfrom models import create_model\nfrom util.visualizer import Visualizer\n\nif __name__ == '__main__':\n opt = TrainOptions().parse()\n\n data_loader = CreateDataLoader(opt)\n dataset = data_loader.load_data()\n dataset_size = len(data_loader)\n print('#training images = %d' % dataset_size)\n\n model = create_model(opt)\n visualizer = Visualizer(opt)\n total_steps = 0\n\n for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):\n epoch_start_time = time.time()\n iter_data_time = time.time()\n epoch_iter = 0\n\n for i, data in enumerate(dataset):\n iter_start_time = time.time()\n if total_steps % opt.print_freq == 0:\n t_data = iter_start_time - iter_data_time\n visualizer.reset()\n total_steps += opt.batchSize\n epoch_iter += opt.batchSize\n model.set_input(data)\n model.optimize_parameters()\n\n if total_steps % opt.display_freq == 0:\n save_result = total_steps % opt.update_html_freq == 0\n visualizer.display_current_results(model.get_current_visuals(), epoch, save_result)\n\n if total_steps % opt.print_freq == 0:\n errors = model.get_current_errors()\n t = (time.time() - iter_start_time) / opt.batchSize\n visualizer.print_current_errors(epoch, epoch_iter, errors, t, t_data)\n if opt.display_id > 0:\n visualizer.plot_current_errors(epoch, float(epoch_iter) / dataset_size, opt, errors)\n\n if total_steps % opt.save_latest_freq == 0:\n print('saving the latest model (epoch %d, total_steps %d)' %\n (epoch, total_steps))\n model.save('latest')\n\n iter_data_time = time.time()\n if epoch % opt.save_epoch_freq == 0:\n print('saving the model at the end of epoch %d, iters %d' %\n (epoch, total_steps))\n model.save('latest')\n model.save(epoch)\n\n print('End of epoch %d / %d \\t Time Taken: %d sec' %\n (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))\n model.update_learning_rate()\n" }, { "path": "drawing_and_style_transfer/util/__init__.py", "content": "" }, { "path": "drawing_and_style_transfer/util/get_data.py", "content": "from __future__ import print_function\nimport os\nimport tarfile\nimport requests\nfrom warnings import warn\nfrom zipfile import ZipFile\nfrom bs4 import BeautifulSoup\nfrom os.path import abspath, isdir, join, basename\n\n\nclass GetData(object):\n \"\"\"\n\n Download CycleGAN or Pix2Pix Data.\n\n Args:\n technique : str\n One of: 'cyclegan' or 'pix2pix'.\n verbose : bool\n If True, print additional information.\n\n Examples:\n >>> from util.get_data import GetData\n >>> gd = GetData(technique='cyclegan')\n >>> new_data_path = gd.get(save_path='./datasets') # options will be displayed.\n\n \"\"\"\n\n def __init__(self, technique='cyclegan', verbose=True):\n url_dict = {\n 'pix2pix': 'https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets',\n 'cyclegan': 'https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets'\n }\n self.url = url_dict.get(technique.lower())\n self._verbose = verbose\n\n def _print(self, text):\n if self._verbose:\n print(text)\n\n @staticmethod\n def _get_options(r):\n soup = BeautifulSoup(r.text, 'lxml')\n options = [h.text for h in soup.find_all('a', href=True)\n if h.text.endswith(('.zip', 'tar.gz'))]\n return options\n\n def _present_options(self):\n r = requests.get(self.url)\n options = self._get_options(r)\n print('Options:\\n')\n for i, o in enumerate(options):\n print(\"{0}: {1}\".format(i, o))\n choice = input(\"\\nPlease enter the number of the \"\n \"dataset above you wish to download:\")\n return options[int(choice)]\n\n def _download_data(self, dataset_url, save_path):\n if not isdir(save_path):\n os.makedirs(save_path)\n\n base = basename(dataset_url)\n temp_save_path = join(save_path, base)\n\n with open(temp_save_path, \"wb\") as f:\n r = requests.get(dataset_url)\n f.write(r.content)\n\n if base.endswith('.tar.gz'):\n obj = tarfile.open(temp_save_path)\n elif base.endswith('.zip'):\n obj = ZipFile(temp_save_path, 'r')\n else:\n raise ValueError(\"Unknown File Type: {0}.\".format(base))\n\n self._print(\"Unpacking Data...\")\n obj.extractall(save_path)\n obj.close()\n os.remove(temp_save_path)\n\n def get(self, save_path, dataset=None):\n \"\"\"\n\n Download a dataset.\n\n Args:\n save_path : str\n A directory to save the data to.\n dataset : str, optional\n A specific dataset to download.\n Note: this must include the file extension.\n If None, options will be presented for you\n to choose from.\n\n Returns:\n save_path_full : str\n The absolute path to the downloaded data.\n\n \"\"\"\n if dataset is None:\n selected_dataset = self._present_options()\n else:\n selected_dataset = dataset\n\n save_path_full = join(save_path, selected_dataset.split('.')[0])\n\n if isdir(save_path_full):\n warn(\"\\n'{0}' already exists. Voiding Download.\".format(\n save_path_full))\n else:\n self._print('Downloading Data...')\n url = \"{0}/{1}\".format(self.url, selected_dataset)\n self._download_data(url, save_path=save_path)\n\n return abspath(save_path_full)\n" }, { "path": "drawing_and_style_transfer/util/html.py", "content": "import dominate\nfrom dominate.tags import *\nimport os\n\n\nclass HTML:\n def __init__(self, web_dir, title, reflesh=0):\n self.title = title\n self.web_dir = web_dir\n self.img_dir = os.path.join(self.web_dir, 'images')\n if not os.path.exists(self.web_dir):\n os.makedirs(self.web_dir)\n if not os.path.exists(self.img_dir):\n os.makedirs(self.img_dir)\n # print(self.img_dir)\n\n self.doc = dominate.document(title=title)\n if reflesh > 0:\n with self.doc.head:\n meta(http_equiv=\"reflesh\", content=str(reflesh))\n\n def get_image_dir(self):\n return self.img_dir\n\n def add_header(self, str):\n with self.doc:\n h3(str)\n\n def add_table(self, border=1):\n self.t = table(border=border, style=\"table-layout: fixed;\")\n self.doc.add(self.t)\n\n def add_images(self, ims, txts, links, width=400):\n self.add_table()\n with self.t:\n with tr():\n for im, txt, link in zip(ims, txts, links):\n with td(style=\"word-wrap: break-word;\", halign=\"center\", valign=\"top\"):\n with p():\n with a(href=os.path.join('images', link)):\n img(style=\"width:%dpx\" % width, src=os.path.join('images', im))\n br()\n p(txt)\n\n def save(self):\n html_file = '%s/index.html' % self.web_dir\n f = open(html_file, 'wt')\n f.write(self.doc.render())\n f.close()\n\n\nif __name__ == '__main__':\n html = HTML('web/', 'test_html')\n html.add_header('hello world')\n\n ims = []\n txts = []\n links = []\n for n in range(4):\n ims.append('image_%d.png' % n)\n txts.append('text_%d' % n)\n links.append('image_%d.png' % n)\n html.add_images(ims, txts, links)\n html.save()\n" }, { "path": "drawing_and_style_transfer/util/image_pool.py", "content": "import random\nimport torch\nfrom torch.autograd import Variable\n\n\nclass ImagePool():\n def __init__(self, pool_size):\n self.pool_size = pool_size\n if self.pool_size > 0:\n self.num_imgs = 0\n self.images = []\n\n def query(self, images):\n if self.pool_size == 0:\n return Variable(images)\n return_images = []\n for image in images:\n image = torch.unsqueeze(image, 0)\n if self.num_imgs < self.pool_size:\n self.num_imgs = self.num_imgs + 1\n self.images.append(image)\n return_images.append(image)\n else:\n p = random.uniform(0, 1)\n if p > 0.5:\n random_id = random.randint(0, self.pool_size - 1)\n tmp = self.images[random_id].clone()\n self.images[random_id] = image\n return_images.append(tmp)\n else:\n return_images.append(image)\n return_images = Variable(torch.cat(return_images, 0))\n return return_images\n" }, { "path": "drawing_and_style_transfer/util/util.py", "content": "from __future__ import print_function\nimport torch\nimport numpy as np\nfrom PIL import Image\nimport os\n\n\n# Converts a Tensor into a Numpy array\n# |imtype|: the desired type of the converted numpy array\ndef tensor2im(image_tensor, imtype=np.uint8):\n image_numpy = image_tensor[0].cpu().float().numpy()\n if image_numpy.shape[0] == 1:\n image_numpy = np.tile(image_numpy, (3, 1, 1))\n image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0\n return image_numpy.astype(imtype)\n\n\ndef diagnose_network(net, name='network'):\n mean = 0.0\n count = 0\n for param in net.parameters():\n if param.grad is not None:\n mean += torch.mean(torch.abs(param.grad.data))\n count += 1\n if count > 0:\n mean = mean / count\n print(name)\n print(mean)\n\n\ndef save_image(image_numpy, image_path):\n image_pil = Image.fromarray(image_numpy)\n image_pil.save(image_path)\n\n\ndef print_numpy(x, val=True, shp=False):\n x = x.astype(np.float64)\n if shp:\n print('shape,', x.shape)\n if val:\n x = x.flatten()\n print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (\n np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))\n\n\ndef mkdirs(paths):\n if isinstance(paths, list) and not isinstance(paths, str):\n for path in paths:\n mkdir(path)\n else:\n mkdir(paths)\n\n\ndef mkdir(path):\n if not os.path.exists(path):\n os.makedirs(path)\n" }, { "path": "drawing_and_style_transfer/util/visualizer.py", "content": "import numpy as np\nimport os\nimport ntpath\nimport time\nfrom . import util\nfrom . import html\nfrom scipy.misc import imresize\n\n\nclass Visualizer():\n def __init__(self, opt):\n # self.opt = opt\n self.display_id = opt.display_id\n self.use_html = opt.isTrain and not opt.no_html\n self.win_size = opt.display_winsize\n self.name = opt.name\n self.opt = opt\n self.saved = False\n if self.display_id > 0:\n import visdom\n self.vis = visdom.Visdom(server=opt.display_server, port=opt.display_port)\n\n if self.use_html:\n self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')\n self.img_dir = os.path.join(self.web_dir, 'images')\n print('create web directory %s...' % self.web_dir)\n util.mkdirs([self.web_dir, self.img_dir])\n self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')\n with open(self.log_name, \"a\") as log_file:\n now = time.strftime(\"%c\")\n log_file.write('================ Training Loss (%s) ================\\n' % now)\n\n def reset(self):\n self.saved = False\n\n # |visuals|: dictionary of images to display or save\n def display_current_results(self, visuals, epoch, save_result):\n if self.display_id > 0: # show images in the browser\n ncols = self.opt.display_single_pane_ncols\n if ncols > 0:\n h, w = next(iter(visuals.values())).shape[:2]\n table_css = \"\"\"\"\"\" % (w, h)\n title = self.name\n label_html = ''\n label_html_row = ''\n nrows = int(np.ceil(len(visuals.items()) / ncols))\n images = []\n idx = 0\n for label, image_numpy in visuals.items():\n label_html_row += '%s' % label\n images.append(image_numpy.transpose([2, 0, 1]))\n idx += 1\n if idx % ncols == 0:\n label_html += '%s' % label_html_row\n label_html_row = ''\n white_image = np.ones_like(image_numpy.transpose([2, 0, 1])) * 255\n while idx % ncols != 0:\n images.append(white_image)\n label_html_row += ''\n idx += 1\n if label_html_row != '':\n label_html += '%s' % label_html_row\n # pane col = image row\n self.vis.images(images, nrow=ncols, win=self.display_id + 1,\n padding=2, opts=dict(title=title + ' images'))\n label_html = '%s
' % label_html\n self.vis.text(table_css + label_html, win=self.display_id + 2,\n opts=dict(title=title + ' labels'))\n else:\n idx = 1\n for label, image_numpy in visuals.items():\n self.vis.image(image_numpy.transpose([2, 0, 1]), opts=dict(title=label),\n win=self.display_id + idx)\n idx += 1\n\n if self.use_html and (save_result or not self.saved): # save images to a html file\n self.saved = True\n for label, image_numpy in visuals.items():\n img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.png' % (epoch, label))\n util.save_image(image_numpy, img_path)\n # update website\n webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, reflesh=1)\n for n in range(epoch, 0, -1):\n webpage.add_header('epoch [%d]' % n)\n ims = []\n txts = []\n links = []\n\n for label, image_numpy in visuals.items():\n img_path = 'epoch%.3d_%s.png' % (n, label)\n ims.append(img_path)\n txts.append(label)\n links.append(img_path)\n webpage.add_images(ims, txts, links, width=self.win_size)\n webpage.save()\n\n # errors: dictionary of error labels and values\n def plot_current_errors(self, epoch, counter_ratio, opt, errors):\n if not hasattr(self, 'plot_data'):\n self.plot_data = {'X': [], 'Y': [], 'legend': list(errors.keys())}\n self.plot_data['X'].append(epoch + counter_ratio)\n self.plot_data['Y'].append([errors[k] for k in self.plot_data['legend']])\n self.vis.line(\n X=np.stack([np.array(self.plot_data['X'])] * len(self.plot_data['legend']), 1),\n Y=np.array(self.plot_data['Y']),\n opts={\n 'title': self.name + ' loss over time',\n 'legend': self.plot_data['legend'],\n 'xlabel': 'epoch',\n 'ylabel': 'loss'},\n win=self.display_id)\n\n # errors: same format as |errors| of plotCurrentErrors\n def print_current_errors(self, epoch, i, errors, t, t_data):\n message = '(epoch: %d, iters: %d, time: %.3f, data: %.3f) ' % (epoch, i, t, t_data)\n for k, v in errors.items():\n message += '%s: %.3f ' % (k, v)\n\n print(message)\n with open(self.log_name, \"a\") as log_file:\n log_file.write('%s\\n' % message)\n\n # save image to the disk\n def save_images(self, webpage, visuals, image_path, aspect_ratio=1.0, index=None, split=1):\n image_dir = webpage.get_image_dir()\n short_path = ntpath.basename(image_path[0])\n name = os.path.splitext(short_path)[0]\n if index is not None:\n name_splits = name.split(\"_\")\n if split == 0:\n name = str(index)\n else:\n name = str(index) + \"_\" + name_splits[split]\n\n webpage.add_header(name)\n ims = []\n txts = []\n links = []\n\n for label, im in visuals.items():\n image_name = '%s_%s.png' % (name, label)\n save_path = os.path.join(image_dir, image_name)\n h, w, _ = im.shape\n if aspect_ratio > 1.0:\n im = imresize(im, (h, int(w * aspect_ratio)), interp='bicubic')\n if aspect_ratio < 1.0:\n im = imresize(im, (int(h / aspect_ratio), w), interp='bicubic')\n util.save_image(im, save_path)\n\n ims.append(image_name)\n txts.append(label)\n links.append(image_name)\n webpage.add_images(ims, txts, links, width=self.win_size)\n" }, { "path": "mnist_to_svhn/data_loader.py", "content": "import torch\nfrom torchvision import datasets\nfrom torchvision import transforms\n\n\ndef get_loader(config):\n \"\"\"Builds and returns Dataloader for MNIST and SVHN dataset.\"\"\"\n\n transform_list = []\n\n if config.use_augmentation:\n transform_list.append(transforms.RandomHorizontalFlip())\n transform_list.append(transforms.RandomRotation(0.1))\n\n transform_list.append(transforms.Scale(config.image_size))\n transform_list.append(transforms.ToTensor())\n transform_list.append(transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)))\n\n transform_test = transforms.Compose([\n transforms.Scale(config.image_size),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\n transform_train = transforms.Compose(transform_list)\n\n svhn = datasets.SVHN(root=config.svhn_path, download=True, transform=transform_train, split='train')\n mnist = datasets.MNIST(root=config.mnist_path, download=True, transform=transform_train, train=True)\n\n svhn_test = datasets.SVHN(root=config.svhn_path, download=True, transform=transform_test, split='test')\n mnist_test = datasets.MNIST(root=config.mnist_path, download=True, transform=transform_test, train=False)\n\n svhn_loader = torch.utils.data.DataLoader(dataset=svhn,\n batch_size=config.batch_size,\n shuffle=config.shuffle,\n num_workers=config.num_workers)\n\n mnist_loader = torch.utils.data.DataLoader(dataset=mnist,\n batch_size=config.batch_size,\n shuffle=config.shuffle,\n num_workers=config.num_workers)\n\n svhn_test_loader = torch.utils.data.DataLoader(dataset=svhn_test,\n batch_size=config.batch_size,\n shuffle=False,\n num_workers=config.num_workers)\n\n mnist_test_loader = torch.utils.data.DataLoader(dataset=mnist_test,\n batch_size=config.batch_size,\n shuffle=False,\n num_workers=config.num_workers)\n\n return svhn_loader, mnist_loader, svhn_test_loader, mnist_test_loader\n" }, { "path": "mnist_to_svhn/download.sh", "content": "mkdir -p mnist\nmkdir -p svhn\n\nwget -O svhn/train_32x32.mat http://ufldl.stanford.edu/housenumbers/train_32x32.mat\nwget -O svhn/test_32x32.mat http://ufldl.stanford.edu/housenumbers/test_32x32.mat\nwget -O svhn/extra_32x32.mat http://ufldl.stanford.edu/housenumbers/extra_32x32.mat" }, { "path": "mnist_to_svhn/main_autoencoder.py", "content": "import argparse\nimport os\nfrom torch.backends import cudnn\n\nfrom solver_autoencoder import Solver\nfrom data_loader import get_loader\n\n\ndef str2bool(v):\n return v.lower() in ('true')\n\n\ndef main(config):\n svhn_loader, mnist_loader, svhn_test_loader, mnist_test_loader = get_loader(config)\n\n solver = Solver(config, svhn_loader, mnist_loader)\n cudnn.benchmark = True\n\n # create directories if not exist\n if not os.path.exists(config.model_path):\n os.makedirs(config.model_path)\n if not os.path.exists(config.sample_path):\n os.makedirs(config.sample_path)\n\n if config.mode == 'train':\n solver.train()\n\n\nif __name__ == '__main__':\n parser = argparse.ArgumentParser()\n\n # model hyper-parameters\n parser.add_argument('--image_size', type=int, default=32)\n parser.add_argument('--g_conv_dim', type=int, default=64)\n parser.add_argument('--d_conv_dim', type=int, default=64)\n parser.add_argument('--num_classes', type=int, default=10)\n\n # training hyper-parameters\n parser.add_argument('--train_iters', type=int, default=15000)\n parser.add_argument('--batch_size', type=int, default=64)\n parser.add_argument('--num_workers', type=int, default=2)\n parser.add_argument('--lr', type=float, default=0.0002)\n parser.add_argument('--beta1', type=float, default=0.5)\n parser.add_argument('--beta2', type=float, default=0.999)\n parser.add_argument('--kl_lambda', type=float, default=0.1)\n\n # misc\n parser.add_argument('--mode', type=str, default='train')\n parser.add_argument('--model_path', type=str, default='./models_autoencoder')\n parser.add_argument('--sample_path', type=str, default='./samples_autoencoder')\n parser.add_argument('--mnist_path', type=str, default='./mnist')\n parser.add_argument('--svhn_path', type=str, default='./svhn')\n parser.add_argument('--log_step', type=int, default=10)\n parser.add_argument('--sample_step', type=int, default=500)\n parser.add_argument('--shuffle', type=bool, default=True)\n parser.add_argument('--use_augmentation', required=True, type=str2bool)\n\n config = parser.parse_args()\n print(config)\n main(config)\n" }, { "path": "mnist_to_svhn/main_mnist_to_svhn.py", "content": "import argparse\nimport logging\nimport os\nfrom torch.backends import cudnn\n\nfrom data_loader import get_loader\nfrom solver_mnist_to_svhn import Solver\n\n\ndef str2bool(v):\n return v.lower() in ('true')\n\n\ndef main(config):\n svhn_loader, mnist_loader, svhn_test_loader, mnist_test_loader = get_loader(config)\n\n solver = Solver(config, svhn_loader, mnist_loader)\n cudnn.benchmark = True\n\n # create directories if not exist\n if not os.path.exists(config.model_path):\n os.makedirs(config.model_path)\n if not os.path.exists(config.sample_path):\n os.makedirs(config.sample_path)\n\n base = config.log_path\n filename = os.path.join(base, str(config.max_items))\n if not os.path.isdir(base):\n os.mkdir(base)\n logging.basicConfig(filename=filename, level=logging.DEBUG)\n\n if config.mode == 'train':\n solver.train()\n elif config.mode == 'sample':\n solver.sample()\n\n\nif __name__ == '__main__':\n parser = argparse.ArgumentParser()\n\n # model hyper-parameters\n parser.add_argument('--image_size', type=int, default=32)\n parser.add_argument('--g_conv_dim', type=int, default=64)\n parser.add_argument('--d_conv_dim', type=int, default=64)\n parser.add_argument('--num_classes', type=int, default=10)\n\n # training hyper-parameters\n parser.add_argument('--train_iters', type=int, default=40000)\n parser.add_argument('--batch_size', type=int, default=1)\n parser.add_argument('--num_workers', type=int, default=2)\n parser.add_argument('--lr', type=float, default=0.0002)\n parser.add_argument('--beta1', type=float, default=0.5)\n parser.add_argument('--beta2', type=float, default=0.999)\n parser.add_argument('--kl_lambda', type=float, default=0.1)\n\n # misc\n parser.add_argument('--mode', type=str, default='train')\n parser.add_argument('--mnist_path', type=str, default='./mnist')\n parser.add_argument('--svhn_path', type=str, default='./svhn')\n parser.add_argument('--log_step', type=int, default=10)\n parser.add_argument('--shuffle', type=bool, default=True)\n\n parser.add_argument('--load_iter', type=int, default=10000)\n parser.add_argument('--sample_step', type=int, default=500)\n parser.add_argument('--num_averaging_runs', type=int, default=1000)\n parser.add_argument('--num_iters_save_model_and_return', type=int, default=5000)\n parser.add_argument('--num_d_iterations', type=int, default=1)\n parser.add_argument('--num_g_iterations', type=int, default=1)\n parser.add_argument('--model_path', type=str, default='./models_ost')\n parser.add_argument('--sample_path', type=str, default='./samples_ost')\n parser.add_argument('--load_path', type=str, default='./models_autoencoder')\n parser.add_argument('--log_path', type=str, default='logs_ost')\n parser.add_argument('--pretrained_g', required=True, type=str2bool)\n parser.add_argument('--save_models_and_samples', required=True, type=str2bool)\n parser.add_argument('--use_augmentation', required=True, type=str2bool)\n parser.add_argument('--one_way_cycle', required=True, type=str2bool)\n parser.add_argument('--freeze_shared', required=True, type=str2bool)\n parser.add_argument('--max_items', type=int, default=1)\n\n config = parser.parse_args()\n print(config)\n main(config)\n" }, { "path": "mnist_to_svhn/main_svhn_to_mnist.py", "content": "import argparse\nimport logging\nimport os\n\nfrom data_loader import get_loader\nfrom solver_svhn_to_mnist import Solver\nfrom torch.backends import cudnn\n\n\ndef str2bool(v):\n return v.lower() in ('true')\n\n\ndef main(config):\n svhn_loader, mnist_loader, svhn_test_loader, mnist_test_loader = get_loader(config)\n\n solver = Solver(config, svhn_loader, mnist_loader)\n cudnn.benchmark = True\n\n # create directories if not exist\n if not os.path.exists(config.model_path):\n os.makedirs(config.model_path)\n if not os.path.exists(config.sample_path):\n os.makedirs(config.sample_path)\n\n base = config.log_path\n filename = os.path.join(base, str(config.max_items))\n if not os.path.isdir(base):\n os.mkdir(base)\n logging.basicConfig(filename=filename, level=logging.DEBUG)\n\n if config.mode == 'train':\n solver.train()\n\n elif config.mode == 'sample':\n solver.sample()\n\n\nif __name__ == '__main__':\n parser = argparse.ArgumentParser()\n\n # model hyper-parameters\n parser.add_argument('--image_size', type=int, default=32)\n parser.add_argument('--g_conv_dim', type=int, default=64)\n parser.add_argument('--d_conv_dim', type=int, default=64)\n parser.add_argument('--num_classes', type=int, default=10)\n\n # training hyper-parameters\n parser.add_argument('--train_iters', type=int, default=40000)\n parser.add_argument('--batch_size', type=int, default=1)\n parser.add_argument('--num_workers', type=int, default=2)\n parser.add_argument('--lr', type=float, default=0.0002)\n parser.add_argument('--beta1', type=float, default=0.5)\n parser.add_argument('--beta2', type=float, default=0.999)\n parser.add_argument('--kl_lambda', type=float, default=0.1)\n\n # misc\n parser.add_argument('--mode', type=str, default='train')\n parser.add_argument('--mnist_path', type=str, default='./mnist')\n parser.add_argument('--svhn_path', type=str, default='./svhn')\n parser.add_argument('--log_step', type=int, default=10)\n parser.add_argument('--shuffle', type=bool, default=True)\n\n parser.add_argument('--load_iter', type=int, default=10000)\n parser.add_argument('--sample_step', type=int, default=500)\n parser.add_argument('--num_averaging_runs', type=int, default=1000)\n parser.add_argument('--num_iters_save_model_and_return', type=int, default=5000)\n parser.add_argument('--num_d_iterations', type=int, default=1)\n parser.add_argument('--num_g_iterations', type=int, default=1)\n parser.add_argument('--model_path', type=str, default='./models_ost')\n parser.add_argument('--sample_path', type=str, default='./samples_ost')\n parser.add_argument('--load_path', type=str, default='./models_autoencoder')\n parser.add_argument('--log_path', type=str, default='logs_ost')\n parser.add_argument('--pretrained_g', required=True, type=str2bool)\n parser.add_argument('--save_models_and_samples', required=True, type=str2bool)\n parser.add_argument('--use_augmentation', required=True, type=str2bool)\n parser.add_argument('--one_way_cycle', required=True, type=str2bool)\n parser.add_argument('--freeze_shared', required=True, type=str2bool)\n parser.add_argument('--max_items', type=int, default=1)\n\n config = parser.parse_args()\n print(config)\n main(config)\n" }, { "path": "mnist_to_svhn/model.py", "content": "import torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef deconv(c_in, c_out, k_size, stride=2, pad=1, bn=True):\n \"\"\"Custom deconvolutional layer for simplicity.\"\"\"\n layers = []\n layers.append(nn.ConvTranspose2d(c_in, c_out, k_size, stride, pad, bias=False))\n if bn:\n layers.append(nn.BatchNorm2d(c_out))\n return nn.Sequential(*layers)\n\n\ndef conv(c_in, c_out, k_size, stride=2, pad=1, bn=True):\n \"\"\"Custom convolutional layer for simplicity.\"\"\"\n layers = []\n layers.append(nn.Conv2d(c_in, c_out, k_size, stride, pad, bias=False))\n if bn:\n layers.append(nn.BatchNorm2d(c_out))\n return nn.Sequential(*layers)\n\n\nclass G11(nn.Module):\n def __init__(self, conv_dim=64):\n super(G11, self).__init__()\n\n # encoding blocks\n self.conv1 = conv(1, conv_dim, 4)\n self.conv1_svhn = conv(3, conv_dim, 4)\n self.conv2 = conv(conv_dim, conv_dim * 2, 4)\n\n # residual blocks\n res_dim = conv_dim * 2\n self.conv3 = conv(res_dim, res_dim, 3, 1, 1)\n self.conv4 = conv(res_dim, res_dim, 3, 1, 1)\n\n # decoding blocks\n self.deconv1 = deconv(conv_dim * 2, conv_dim, 4)\n self.deconv2 = deconv(conv_dim, 1, 4, bn=False)\n self.deconv2_svhn = deconv(conv_dim, 3, 4, bn=False)\n\n def forward(self, x, svhn=False):\n if svhn:\n out = F.leaky_relu(self.conv1_svhn(x), 0.05) # (?, 64, 16, 16)\n else:\n out = F.leaky_relu(self.conv1(x), 0.05) # (?, 64, 16, 16)\n\n out = F.leaky_relu(self.conv2(out), 0.05) # (?, 128, 8, 8)\n out = F.leaky_relu(self.conv3(out), 0.05) # ( \" )\n out = F.leaky_relu(self.conv4(out), 0.05) # ( \" )\n out = F.leaky_relu(self.deconv1(out), 0.05) # (?, 64, 16, 16)\n\n if svhn:\n out = F.tanh(self.deconv2_svhn(out)) # (?, 3, 32, 32)\n else:\n out = F.tanh(self.deconv2(out)) # (?, 3, 32, 32)\n\n return out\n\n def encode(self, x, svhn=False):\n\n if svhn:\n out = F.leaky_relu(self.conv1_svhn(x), 0.05) # (?, 64, 16, 16)\n else:\n out = F.leaky_relu(self.conv1(x), 0.05) # (?, 64, 16, 16)\n\n out = F.leaky_relu(self.conv2(out), 0.05) # (?, 128, 8, 8)\n out = F.leaky_relu(self.conv3(out), 0.05) # ( \" )\n\n return out\n\n def decode(self, out, svhn=False):\n\n out = F.leaky_relu(self.conv4(out), 0.05)\n out = F.leaky_relu(self.deconv1(out), 0.05) # (?, 64, 16, 16)\n\n if svhn:\n out = F.tanh(self.deconv2_svhn(out)) # (?, 3, 32, 32)\n else:\n out = F.tanh(self.deconv2(out)) # (?, 3, 32, 32)\n\n return out\n\n def encode_params(self):\n layers_basic = list(self.conv1_svhn.parameters()) + \\\n list(self.conv1.parameters())\n layers_basic += list(self.conv2.parameters())\n layers_basic += list(self.conv3.parameters())\n\n return layers_basic\n\n def decode_params(self):\n layers_basic = list(self.deconv2_svhn.parameters()) + \\\n list(self.deconv2.parameters())\n layers_basic += list(self.deconv1.parameters())\n layers_basic += list(self.conv4.parameters())\n\n return layers_basic\n\n def unshared_parameters(self):\n return list(self.deconv2_svhn.parameters()) + list(self.conv1_svhn.parameters()) + \\\n list(self.deconv2.parameters()) + list(self.conv1.parameters())\n\n\nclass G22(nn.Module):\n def __init__(self, conv_dim=64):\n super(G22, self).__init__()\n\n # encoding blocks\n self.conv1 = conv(3, conv_dim, 4)\n self.conv1_mnist = conv(1, conv_dim, 4)\n self.conv2 = conv(conv_dim, conv_dim * 2, 4)\n\n # residual blocks\n res_dim = conv_dim * 2\n self.conv3 = conv(res_dim, res_dim, 3, 1, 1)\n self.conv4 = conv(res_dim, res_dim, 3, 1, 1)\n\n # decoding blocks\n self.deconv1 = deconv(conv_dim * 2, conv_dim, 4)\n self.deconv2 = deconv(conv_dim, 3, 4, bn=False)\n self.deconv2_mnist = deconv(conv_dim, 1, 4, bn=False)\n\n def forward(self, x, mnist=False):\n if mnist:\n out = F.leaky_relu(self.conv1_mnist(x), 0.05) # (?, 64, 16, 16)\n else:\n out = F.leaky_relu(self.conv1(x), 0.05) # (?, 64, 16, 16)\n\n out = F.leaky_relu(self.conv2(out), 0.05) # (?, 128, 8, 8)\n out = F.leaky_relu(self.conv3(out), 0.05) # ( \" )\n out = F.leaky_relu(self.conv4(out), 0.05) # ( \" )\n out = F.leaky_relu(self.deconv1(out), 0.05) # (?, 64, 16, 16)\n\n if mnist:\n out = F.tanh(self.deconv2_mnist(out)) # (?, 3, 32, 32)\n else:\n out = F.tanh(self.deconv2(out)) # (?, 3, 32, 32)\n\n return out\n\n def encode(self, x, mnist=False):\n\n if mnist:\n out = F.leaky_relu(self.conv1_mnist(x), 0.05) # (?, 64, 16, 16)\n else:\n out = F.leaky_relu(self.conv1(x), 0.05) # (?, 64, 16, 16)\n\n out = F.leaky_relu(self.conv2(out), 0.05) # (?, 128, 8, 8)\n out = F.leaky_relu(self.conv3(out), 0.05) # ( \" )\n\n return out\n\n def decode(self, out, mnist=False):\n\n out = F.leaky_relu(self.conv4(out), 0.05)\n out = F.leaky_relu(self.deconv1(out), 0.05) # (?, 64, 16, 16)\n\n if mnist:\n out = F.tanh(self.deconv2_mnist(out)) # (?, 3, 32, 32)\n else:\n out = F.tanh(self.deconv2(out)) # (?, 3, 32, 32)\n\n return out\n\n def encode_params(self):\n layers_basic = list(self.conv1_mnist.parameters()) + \\\n list(self.conv1.parameters())\n layers_basic += list(self.conv2.parameters())\n layers_basic += list(self.conv3.parameters())\n\n return layers_basic\n\n def decode_params(self):\n layers_basic = list(self.deconv2_mnist.parameters()) + \\\n list(self.deconv2.parameters())\n layers_basic += list(self.deconv1.parameters())\n layers_basic += list(self.conv4.parameters())\n\n return layers_basic\n\n def unshared_parameters(self):\n return list(self.deconv2_mnist.parameters()) + list(self.conv1_mnist.parameters()) + \\\n list(self.deconv2.parameters()) + list(self.conv1.parameters())\n\n\nclass D1(nn.Module):\n \"\"\"Discriminator for mnist.\"\"\"\n\n def __init__(self, conv_dim=64, use_labels=False):\n super(D1, self).__init__()\n self.conv1 = conv(1, conv_dim, 4, bn=False)\n self.conv2 = conv(conv_dim, conv_dim * 2, 4)\n self.conv3 = conv(conv_dim * 2, conv_dim * 4, 4)\n n_out = 11 if use_labels else 1\n self.fc = conv(conv_dim * 4, n_out, 4, 1, 0, False)\n\n def forward(self, x_0):\n out = F.leaky_relu(self.conv1(x_0), 0.05) # (?, 64, 16, 16)\n out = F.leaky_relu(self.conv2(out), 0.05) # (?, 128, 8, 8)\n out = F.leaky_relu(self.conv3(out), 0.05) # (?, 256, 4, 4)\n out_0 = self.fc(out).squeeze()\n\n return out_0\n\n\nclass D2(nn.Module):\n \"\"\"Discriminator for svhn.\"\"\"\n\n def __init__(self, conv_dim=64, use_labels=False):\n super(D2, self).__init__()\n self.conv1 = conv(3, conv_dim, 4, bn=False)\n self.conv2 = conv(conv_dim, conv_dim * 2, 4)\n self.conv3 = conv(conv_dim * 2, conv_dim * 4, 4)\n n_out = 11 if use_labels else 1\n self.fc = conv(conv_dim * 4, n_out, 4, 1, 0, False)\n\n def forward(self, x_0):\n out = F.leaky_relu(self.conv1(x_0), 0.05) # (?, 64, 16, 16)\n out = F.leaky_relu(self.conv2(out), 0.05) # (?, 128, 8, 8)\n out = F.leaky_relu(self.conv3(out), 0.05) # (?, 256, 4, 4)\n out_0 = self.fc(out).squeeze()\n\n return out_0\n" }, { "path": "mnist_to_svhn/solver_autoencoder.py", "content": "import os\n\nimport numpy as np\nimport scipy.io\nimport torch\nfrom torch import optim\nfrom torch.autograd import Variable\n\nfrom model import D1, D2\nfrom model import G11, G22\n\n\nclass Solver(object):\n def __init__(self, config, svhn_loader, mnist_loader):\n self.svhn_loader = svhn_loader\n self.mnist_loader = mnist_loader\n self.g11 = None\n self.g22 = None\n self.d1 = None\n self.d2 = None\n self.g_optimizer = None\n self.d_optimizer = None\n self.num_classes = config.num_classes\n self.beta1 = config.beta1\n self.beta2 = config.beta2\n self.g_conv_dim = config.g_conv_dim\n self.d_conv_dim = config.d_conv_dim\n self.train_iters = config.train_iters\n self.batch_size = config.batch_size\n self.lr = config.lr\n self.log_step = config.log_step\n self.sample_step = config.sample_step\n self.sample_path = config.sample_path\n self.model_path = config.model_path\n self.kl_lambda = config.kl_lambda\n self.build_model()\n\n def build_model(self):\n \"\"\"Builds a generator and a discriminator.\"\"\"\n self.g11 = G11(conv_dim=self.g_conv_dim)\n self.g22 = G22(conv_dim=self.g_conv_dim)\n self.d1 = D1(conv_dim=self.d_conv_dim, use_labels=False)\n self.d2 = D2(conv_dim=self.d_conv_dim, use_labels=False)\n\n g_params = list(self.g11.parameters()) + list(self.g22.parameters())\n d_params = list(self.d1.parameters()) + list(self.d2.parameters())\n\n self.g_optimizer = optim.Adam(g_params, self.lr, [self.beta1, self.beta2])\n self.d_optimizer = optim.Adam(d_params, self.lr, [self.beta1, self.beta2])\n\n if torch.cuda.is_available():\n self.g11.cuda()\n self.g22.cuda()\n self.d1.cuda()\n self.d2.cuda()\n\n def merge_images(self, sources, targets, k=10):\n _, _, h, w = sources.shape\n row = int(np.sqrt(self.batch_size))\n merged = np.zeros([3, row * h, row * w * 2])\n for idx, (s, t) in enumerate(zip(sources, targets)):\n i = idx // row\n j = idx % row\n merged[:, i * h:(i + 1) * h, (j * 2) * h:(j * 2 + 1) * h] = s\n merged[:, i * h:(i + 1) * h, (j * 2 + 1) * h:(j * 2 + 2) * h] = t\n return merged.transpose(1, 2, 0)\n\n def to_var(self, x):\n \"\"\"Converts numpy to variable.\"\"\"\n if torch.cuda.is_available():\n x = x.cuda()\n return Variable(x)\n\n def to_data(self, x):\n \"\"\"Converts variable to numpy.\"\"\"\n if torch.cuda.is_available():\n x = x.cpu()\n return x.data.numpy()\n\n def reset_grad(self):\n \"\"\"Zeros the gradient buffers.\"\"\"\n self.g_optimizer.zero_grad()\n self.d_optimizer.zero_grad()\n\n def _compute_kl(self, mu):\n mu_2 = torch.pow(mu, 2)\n encoding_loss = torch.mean(mu_2)\n return encoding_loss\n\n def train(self):\n svhn_iter = iter(self.svhn_loader)\n mnist_iter = iter(self.mnist_loader)\n iter_per_epoch = min(len(svhn_iter), len(mnist_iter))\n\n # fixed mnist and svhn for sampling\n fixed_svhn = self.to_var(svhn_iter.next()[0])\n fixed_mnist = self.to_var(mnist_iter.next()[0])\n\n # Train autoencoder for mnist\n for step in range(self.train_iters + 1):\n # reset data_iter for each epoch\n if (step + 1) % iter_per_epoch == 0:\n mnist_iter = iter(self.mnist_loader)\n\n # mnist dataset\n mnist_data, m_labels_data = mnist_iter.next()\n mnist, m_labels = self.to_var(mnist_data), self.to_var(m_labels_data)\n\n # ============ train D ============#\n # train with real images\n self.reset_grad()\n \n out = self.d1(mnist)\n d1_loss = torch.mean((out - 1) ** 2)\n\n d_mnist_loss = d1_loss\n d_real_loss = d1_loss\n d_real_loss.backward()\n self.d_optimizer.step()\n\n # train with fake images\n self.reset_grad()\n fake_mnist = self.g22.forward(mnist, mnist=True)\n out = self.d1(fake_mnist)\n d2_loss = torch.mean(out ** 2)\n\n d_fake_loss = d2_loss\n d_fake_loss.backward()\n self.d_optimizer.step()\n\n # ============ train G ============\n self.reset_grad()\n fake_mnist = self.g22.forward(mnist, mnist=True)\n out = self.d1(fake_mnist)\n g_loss = torch.mean((out - 1) ** 2)\n g_loss += torch.mean((mnist - fake_mnist) ** 2)\n em = self.g22.encode(mnist, mnist=True)\n g_loss += self.kl_lambda * self._compute_kl(em)\n\n g_loss.backward()\n self.g_optimizer.step()\n\n # print the log info\n if (step + 1) % self.log_step == 0:\n print('Step [%d/%d], d_real_loss: %.4f, d_mnist_loss: %.4f, d_svhn_loss: %.4f, '\n 'g_loss: %.4f'\n % (step + 1, self.train_iters, d_real_loss.data[0], d_mnist_loss.data[0],\n d_fake_loss.data[0], g_loss.data[0]))\n\n # save the sampled images\n if (step + 1) % self.sample_step == 0:\n fake_mnist = self.g22.forward(fixed_mnist, mnist=True)\n mnist, fake_mnist = self.to_data(fixed_mnist), self.to_data(fake_mnist)\n\n merged = self.merge_images(mnist, fake_mnist)\n path = os.path.join(self.sample_path, 'sample-%d-m-s.png' % (step + 1))\n scipy.misc.imsave(path, merged)\n print('saved %s' % path)\n\n if (step + 1) % 10000 == 0:\n # save the model parameters for each epoch\n g22_path = os.path.join(self.model_path, 'g22-%d.pkl' % (step + 1))\n d1_path = os.path.join(self.model_path, 'd1-%d.pkl' % (step + 1))\n torch.save(self.g22.state_dict(), g22_path)\n torch.save(self.d1.state_dict(), d1_path)\n\n # Train autoencoder for svhn\n for step in range(self.train_iters + 1):\n # reset data_iter for each epoch\n if (step + 1) % iter_per_epoch == 0:\n svhn_iter = iter(self.svhn_loader)\n\n # load svhn and mnist dataset\n svhn_data, s_labels_data = svhn_iter.next()\n svhn, s_labels = self.to_var(svhn_data), self.to_var(s_labels_data).long().squeeze()\n\n # ============ train D ============#\n\n # train with real images\n self.reset_grad()\n\n out = self.d2(svhn)\n d2_loss = torch.mean((out - 1) ** 2)\n\n d_svhn_loss = d2_loss\n d_real_loss = d2_loss\n d_real_loss.backward()\n self.d_optimizer.step()\n\n # train with fake images\n self.reset_grad()\n\n fake_svhn = self.g11.forward(svhn, svhn=True)\n out = self.d2(fake_svhn)\n d1_loss = torch.mean(out ** 2)\n\n d_fake_loss = d1_loss\n d_fake_loss.backward()\n self.d_optimizer.step()\n\n # ============ train G ============#\n\n self.reset_grad()\n fake_svhn = self.g11.forward(svhn, svhn=True)\n out = self.d2(fake_svhn)\n g_loss = torch.mean((out - 1) ** 2)\n g_loss += torch.mean((svhn - fake_svhn) ** 2)\n es = self.g11.encode(svhn, svhn=True)\n g_loss += self.kl_lambda * self._compute_kl(es)\n\n g_loss.backward()\n self.g_optimizer.step()\n\n # print the log info\n if (step + 1) % self.log_step == 0:\n print('Step [%d/%d], d_real_loss: %.4f, d_mnist_loss: %.4f,'\n 'd_fake_loss: %.4f, g_loss: %.4f'\n % (step + 1, self.train_iters, d_real_loss.data[0],\n d_svhn_loss.data[0], d_fake_loss.data[0], g_loss.data[0]))\n\n # save the sampled images\n if (step + 1) % self.sample_step == 0:\n fake_svhn = self.g11.forward(fixed_svhn, svhn=True)\n svhn, fake_svhn = self.to_data(fixed_svhn), self.to_data(fake_svhn)\n\n merged = self.merge_images(svhn, fake_svhn)\n path = os.path.join(self.sample_path, 'sample-%d-s-m.png' % (step + 1))\n scipy.misc.imsave(path, merged)\n print('saved %s' % path)\n\n if (step + 1) % 10000 == 0:\n # save the model parameters for each epoch\n g11_path = os.path.join(self.model_path, 'g11-%d.pkl' % (step + 1))\n d2_path = os.path.join(self.model_path, 'd2-%d.pkl' % (step + 1))\n torch.save(self.g11.state_dict(), g11_path)\n torch.save(self.d2.state_dict(), d2_path)\n" }, { "path": "mnist_to_svhn/solver_mnist_to_svhn.py", "content": "import os\n\nimport numpy as np\nimport scipy.io\nimport torch\nfrom torch import optim\nfrom torch.autograd import Variable\n\nfrom model import D1, D2\nfrom model import G11\n\n\nclass Solver(object):\n def __init__(self, config, svhn_loader, mnist_loader):\n self.config = config\n self.svhn_loader = svhn_loader\n self.mnist_loader = mnist_loader\n self.g11 = None\n self.g22 = None\n self.d1 = None\n self.d2 = None\n self.g_optimizer = None\n self.num_classes = config.num_classes\n self.beta1 = config.beta1\n self.beta2 = config.beta2\n self.g_conv_dim = config.g_conv_dim\n self.d_conv_dim = config.d_conv_dim\n self.train_iters = config.train_iters\n self.batch_size = config.batch_size\n self.lr = config.lr\n self.kl_lambda = config.kl_lambda\n self.log_step = config.log_step\n self.sample_step = config.sample_step\n self.sample_path = config.sample_path\n self.model_path = config.model_path\n self.g11_load_path = os.path.join(config.load_path, \"g11-\" + str(config.load_iter) + \".pkl\")\n self.d1_load_path = os.path.join(config.load_path, \"d1-\" + str(config.load_iter) + \".pkl\")\n self.g22_load_path = os.path.join(config.load_path, \"g22-\" + str(config.load_iter) + \".pkl\")\n self.d2_load_path = os.path.join(config.load_path, \"d2-\" + str(config.load_iter) + \".pkl\")\n self.build_model()\n\n def build_model(self):\n \"\"\"Builds a generator and a discriminator.\"\"\"\n self.g11 = G11(conv_dim=self.g_conv_dim)\n self.g_optimizer = optim.Adam(list(self.g11.encode_params()) + list(self.g11.decode_params()), self.lr,\n [self.beta1, self.beta2])\n self.unshared_optimizer = optim.Adam(list(self.g11.unshared_parameters()), self.lr,\n [self.beta1, self.beta2])\n\n self.d1 = D1(conv_dim=self.d_conv_dim, use_labels=False)\n self.d2 = D2(conv_dim=self.d_conv_dim, use_labels=False)\n\n self.d_optimizer = optim.Adam(list(self.d1.parameters()) + list(self.d2.parameters()), self.lr,\n [self.beta1, self.beta2])\n\n if torch.cuda.is_available():\n self.g11.cuda()\n self.d1.cuda()\n self.d2.cuda()\n\n def merge_images(self, sources, targets, k=10):\n _, _, h, w = sources.shape\n row = int(np.sqrt(self.batch_size)) + 1\n merged = np.zeros([3, row * h, row * w * 2])\n for idx, (s, t) in enumerate(zip(sources, targets)):\n i = idx // row\n j = idx % row\n merged[:, i * h:(i + 1) * h, (j * 2) * h:(j * 2 + 1) * h] = s\n merged[:, i * h:(i + 1) * h, (j * 2 + 1) * h:(j * 2 + 2) * h] = t\n return merged.transpose(1, 2, 0)\n\n def to_var(self, x, volatile=False):\n \"\"\"Converts numpy to variable.\"\"\"\n if torch.cuda.is_available():\n x = x.cuda()\n if volatile:\n return Variable(x, volatile=True)\n return Variable(x)\n\n def to_no_grad_var(self, x):\n x = self.to_data(x, no_numpy=True)\n return self.to_var(x, volatile=True)\n\n def to_data(self, x, no_numpy=False):\n \"\"\"Converts variable to numpy.\"\"\"\n if torch.cuda.is_available():\n x = x.cpu()\n if no_numpy:\n return x.data\n return x.data.numpy()\n\n def reset_grad(self):\n \"\"\"Zeros the gradient buffers.\"\"\"\n self.unshared_optimizer.zero_grad()\n self.g_optimizer.zero_grad()\n self.d_optimizer.zero_grad()\n\n def _compute_kl(self, mu):\n\n mu_2 = torch.pow(mu, 2)\n encoding_loss = torch.mean(mu_2)\n return encoding_loss\n\n def train(self):\n self.build_model()\n if self.config.pretrained_g:\n self.g11.load_state_dict(torch.load(self.g11_load_path))\n\n svhn_iter = iter(self.svhn_loader)\n mnist_iter = iter(self.mnist_loader)\n iter_per_epoch = min(len(svhn_iter), len(mnist_iter))\n\n # fixed mnist and svhn for sampling\n svhn_fixed_data, svhn_fixed_labels = svhn_iter.next()\n mnist_fixed_data, mnist_fixed_labels = mnist_iter.next()\n fixed_mnist = self.to_var(mnist_fixed_data)\n counter = 0\n\n for step in range(self.train_iters + 1):\n\n # reset data_iter for each epoch\n if (step + 1) % iter_per_epoch == 0:\n mnist_iter = iter(self.mnist_loader)\n svhn_iter = iter(self.svhn_loader)\n\n # load svhn and mnist dataset\n svhn_data, s_labels_data = svhn_iter.next()\n mnist_data, m_labels_data = mnist_iter.next()\n svhn, s_labels = self.to_var(svhn_data), self.to_var(s_labels_data).long().squeeze()\n mnist, m_labels = self.to_var(mnist_data), self.to_var(m_labels_data)\n\n # This sets the maximum number of items for A domain\n # We assume max_items is a multiple of batch_size\n # And reset mnist loader when we pass the number of allowed items.\n if self.batch_size > self.config.max_items:\n exit(-1)\n elif self.batch_size == self.config.max_items:\n mnist = fixed_mnist\n elif self.batch_size < self.config.max_items:\n counter += 1\n if counter * self.batch_size >= self.config.max_items:\n mnist_iter = iter(self.mnist_loader)\n counter = 0\n\n # ============ train D ============#\n # train with real images\n self.reset_grad()\n out = self.d1(mnist)\n d1_loss = torch.mean((out - 1) ** 2)\n\n out = self.d2(svhn)\n d2_loss = torch.mean((out - 1) ** 2)\n\n d_mnist_loss = d1_loss\n d_svhn_loss = d2_loss\n # Only optimizing d1\n d_real_loss = d1_loss + d2_loss\n d_real_loss.backward()\n self.d_optimizer.step()\n\n # train with fake images\n self.reset_grad()\n es = self.g11.encode(svhn, svhn=True)\n fake_mnist = self.g11.decode(es)\n out = self.d1(fake_mnist)\n d2_loss = torch.mean(out ** 2)\n\n em = self.g11.encode(mnist)\n fake_svhn = self.g11.decode(em, svhn=True)\n out = self.d2(fake_svhn)\n d1_loss = torch.mean(out ** 2)\n\n d_fake_loss = d2_loss + d1_loss\n d_fake_loss.backward()\n self.d_optimizer.step()\n\n # ============ train G ============#\n\n # train mnist-svhn-mnist cycle\n self.reset_grad()\n es = self.g11.encode(svhn, svhn=True)\n fake_mnist = self.g11.decode(es)\n out = self.d1(fake_mnist)\n g_loss = torch.mean((out - 1) ** 2)\n\n em = self.g11.encode(mnist)\n fake_svhn = self.g11.decode(em, svhn=True)\n out = self.d2(fake_svhn)\n g_loss += torch.mean((out - 1) ** 2)\n\n self.reset_grad()\n em = self.g11.encode(mnist)\n fake_mnist = self.g11.decode(em)\n g_loss += torch.mean((mnist - fake_mnist) ** 2)\n\n if self.config.one_way_cycle:\n em = self.g11.encode(mnist)\n fake_svhn = self.g11.decode(em, svhn=True)\n es = self.g11.encode(fake_svhn, svhn=True)\n fake_mnist = self.g11.decode(es)\n g_loss += torch.mean((mnist - fake_mnist) ** 2)\n\n g_loss.backward()\n self.unshared_optimizer.step()\n\n if not self.config.freeze_shared:\n self.reset_grad()\n es = self.g11.encode(svhn, svhn=True)\n fake_es = self.g11.decode(es, svhn=True)\n g_loss = torch.mean((svhn - fake_es) ** 2)\n g_loss += self.kl_lambda * self._compute_kl(es)\n\n g_loss.backward()\n self.g_optimizer.step()\n\n # print the log info\n if (step + 1) % self.log_step == 0:\n print('Step [%d/%d], d_real_loss: %.4f, d_mnist_loss: %.4f, d_svhn_loss: %.4f, '\n 'd_fake_loss: %.4f, g_loss: %.4f'\n % (step + 1, self.train_iters, d_real_loss.data[0], d_mnist_loss.data[0],\n d_svhn_loss.data[0], d_fake_loss.data[0], g_loss.data[0]))\n\n # save the sampled images\n if (step + 1) % self.sample_step == 0:\n em = self.g11.encode(fixed_mnist)\n fake_svhn_var = self.g11.decode(em, svhn=True)\n fake_svhn = self.to_data(fake_svhn_var)\n if self.config.save_models_and_samples:\n merged = self.merge_images(mnist_fixed_data, fake_svhn)\n path = os.path.join(self.sample_path, 'sample-%d-s-m.png' % (step + 1))\n scipy.misc.imsave(path, merged)\n print('saved %s' % path)\n\n if (step + 1) % self.config.num_iters_save_model_and_return == 0:\n # save the model parameters for each epoch\n if self.config.save_models_and_samples:\n g11_path = os.path.join(self.model_path, 'g11-%d.pkl' % (step + 1))\n d1_path = os.path.join(self.model_path, 'd1-%d.pkl' % (step + 1))\n d2_path = os.path.join(self.model_path, 'd2-%d.pkl' % (step + 1))\n torch.save(self.g11.state_dict(), g11_path)\n torch.save(self.d1.state_dict(), d1_path)\n torch.save(self.d2.state_dict(), d2_path)\n\n return\n" }, { "path": "mnist_to_svhn/solver_svhn_to_mnist.py", "content": "import os\n\nimport numpy as np\nimport scipy.io\nimport torch\nfrom torch import optim\nfrom torch.autograd import Variable\n\nfrom model import D1, D2\nfrom model import G22\n\n\nclass Solver(object):\n def __init__(self, config, svhn_loader, mnist_loader):\n self.config = config\n self.svhn_loader = svhn_loader\n self.mnist_loader = mnist_loader\n self.g11 = None\n self.g22 = None\n self.d1 = None\n self.d2 = None\n self.g_optimizer = None\n self.num_classes = config.num_classes\n self.beta1 = config.beta1\n self.beta2 = config.beta2\n self.g_conv_dim = config.g_conv_dim\n self.d_conv_dim = config.d_conv_dim\n self.train_iters = config.train_iters\n self.batch_size = config.batch_size\n self.lr = config.lr\n self.kl_lambda = config.kl_lambda\n self.log_step = config.log_step\n self.sample_step = config.sample_step\n self.sample_path = config.sample_path\n self.model_path = config.model_path\n self.g11_load_path = os.path.join(config.load_path, \"g11-\" + str(config.load_iter) + \".pkl\")\n self.d1_load_path = os.path.join(config.load_path, \"d1-\" + str(config.load_iter) + \".pkl\")\n self.g22_load_path = os.path.join(config.load_path, \"g22-\" + str(config.load_iter) + \".pkl\")\n self.d2_load_path = os.path.join(config.load_path, \"d2-\" + str(config.load_iter) + \".pkl\")\n self.build_model()\n\n def build_model(self):\n \"\"\"Builds a generator and a discriminator.\"\"\"\n self.g22 = G22(conv_dim=self.g_conv_dim)\n self.g_optimizer = optim.Adam(list(self.g22.encode_params()) + list(self.g22.decode_params()), self.lr,\n [self.beta1, self.beta2])\n self.unshared_optimizer = optim.Adam(list(self.g22.unshared_parameters()), self.lr,\n [self.beta1, self.beta2])\n\n self.d1 = D1(conv_dim=self.d_conv_dim, use_labels=False)\n self.d2 = D2(conv_dim=self.d_conv_dim, use_labels=False)\n\n self.d_optimizer = optim.Adam(list(self.d1.parameters()) + list(self.d2.parameters()), self.lr,\n [self.beta1, self.beta2])\n\n if torch.cuda.is_available():\n self.g22.cuda()\n self.d1.cuda()\n self.d2.cuda()\n\n def merge_images(self, sources, targets, k=10):\n _, _, h, w = sources.shape\n row = int(np.sqrt(self.batch_size)) + 1\n merged = np.zeros([3, row * h, row * w * 2])\n for idx, (s, t) in enumerate(zip(sources, targets)):\n i = idx // row\n j = idx % row\n merged[:, i * h:(i + 1) * h, (j * 2) * h:(j * 2 + 1) * h] = s\n merged[:, i * h:(i + 1) * h, (j * 2 + 1) * h:(j * 2 + 2) * h] = t\n return merged.transpose(1, 2, 0)\n\n def to_var(self, x, volatile=False):\n \"\"\"Converts numpy to variable.\"\"\"\n if torch.cuda.is_available():\n x = x.cuda()\n if volatile:\n return Variable(x, volatile=True)\n return Variable(x)\n\n def to_no_grad_var(self, x):\n x = self.to_data(x, no_numpy=True)\n return self.to_var(x, volatile=True)\n\n def to_data(self, x, no_numpy=False):\n \"\"\"Converts variable to numpy.\"\"\"\n if torch.cuda.is_available():\n x = x.cpu()\n if no_numpy:\n return x.data\n return x.data.numpy()\n\n def reset_grad(self):\n \"\"\"Zeros the gradient buffers.\"\"\"\n self.unshared_optimizer.zero_grad()\n self.g_optimizer.zero_grad()\n self.d_optimizer.zero_grad()\n\n def _compute_kl(self, mu):\n mu_2 = torch.pow(mu, 2)\n encoding_loss = torch.mean(mu_2)\n return encoding_loss\n\n def train(self):\n self.build_model()\n if self.config.pretrained_g:\n self.g22.load_state_dict(torch.load(self.g22_load_path))\n\n svhn_iter = iter(self.svhn_loader)\n mnist_iter = iter(self.mnist_loader)\n iter_per_epoch = min(len(svhn_iter), len(mnist_iter))\n\n # fixed mnist and svhn for sampling\n svhn_fixed_data, svhn_fixed_labels = svhn_iter.next()\n mnist_fixed_data, mnist_fixed_labels = mnist_iter.next()\n fixed_svhn = self.to_var(svhn_fixed_data)\n counter = 0\n\n for step in range(self.train_iters + 1):\n\n # reset data_iter for each epoch\n if (step + 1) % iter_per_epoch == 0:\n mnist_iter = iter(self.mnist_loader)\n svhn_iter = iter(self.svhn_loader)\n\n # load svhn and mnist dataset\n svhn_data, s_labels_data = svhn_iter.next()\n mnist_data, m_labels_data = mnist_iter.next()\n svhn, s_labels = self.to_var(svhn_data), self.to_var(s_labels_data).long().squeeze()\n mnist, m_labels = self.to_var(mnist_data), self.to_var(m_labels_data)\n\n # This sets the maximum number of items for A domain\n # We assume max_items is a multiple of batch_size\n # And reset mnist loader when we pass the number of allowed items.\n if self.batch_size > self.config.max_items:\n exit(-1)\n elif self.batch_size == self.config.max_items:\n svhn = fixed_svhn\n elif self.batch_size < self.config.max_items:\n counter += 1\n if counter * self.batch_size >= self.config.max_items:\n svhn_iter = iter(self.svhn_loader)\n counter = 0\n\n # ============ train D ============#\n # train with real images\n self.reset_grad()\n out = self.d1(mnist)\n d1_loss = torch.mean((out - 1) ** 2)\n\n out = self.d2(svhn)\n d2_loss = torch.mean((out - 1) ** 2)\n\n d_mnist_loss = d1_loss\n d_svhn_loss = d2_loss\n # Only optimizing d1\n d_real_loss = d1_loss + d2_loss\n d_real_loss.backward()\n self.d_optimizer.step()\n\n # train with fake images\n self.reset_grad()\n es = self.g22.encode(svhn)\n fake_mnist = self.g22.decode(es, mnist=True)\n out = self.d1(fake_mnist)\n d2_loss = torch.mean(out ** 2)\n\n em = self.g22.encode(mnist, mnist=True)\n fake_svhn = self.g22.decode(em)\n out = self.d2(fake_svhn)\n d1_loss = torch.mean(out ** 2)\n\n d_fake_loss = d2_loss + d1_loss\n d_fake_loss.backward()\n self.d_optimizer.step()\n\n # ============ train G ============#\n\n self.reset_grad()\n es = self.g22.encode(svhn)\n fake_mnist = self.g22.decode(es, mnist=True)\n out = self.d1(fake_mnist)\n g_loss = torch.mean((out - 1) ** 2)\n\n em = self.g22.encode(mnist, mnist=True)\n fake_svhn = self.g22.decode(em)\n out = self.d2(fake_svhn)\n g_loss += torch.mean((out - 1) ** 2)\n\n self.reset_grad()\n es = self.g22.encode(svhn)\n fake_svhn = self.g22.decode(es)\n g_loss += torch.mean((svhn - fake_svhn) ** 2)\n\n if self.config.one_way_cycle:\n es = self.g22.encode(svhn)\n fake_mnist = self.g22.decode(es, mnist=True)\n es = self.g22.encode(fake_mnist, mnist=True)\n fake_svhn = self.g22.decode(es)\n g_loss += torch.mean((svhn - fake_svhn) ** 2)\n\n g_loss.backward()\n self.unshared_optimizer.step()\n\n if not self.config.freeze_shared:\n self.reset_grad()\n em = self.g22.encode(mnist, mnist=True)\n fake_em = self.g22.decode(em, mnist=True)\n g_loss = torch.mean((mnist - fake_em) ** 2)\n g_loss += self.kl_lambda * self._compute_kl(em)\n\n g_loss.backward()\n self.g_optimizer.step()\n\n # print the log info\n if (step + 1) % self.log_step == 0:\n print('Step [%d/%d], d_real_loss: %.4f, d_mnist_loss: %.4f, d_svhn_loss: %.4f, '\n 'd_fake_loss: %.4f, g_loss: %.4f'\n % (step + 1, self.train_iters, d_real_loss.data[0], d_mnist_loss.data[0],\n d_svhn_loss.data[0], d_fake_loss.data[0], g_loss.data[0]))\n\n # save the sampled images\n if (step + 1) % self.sample_step == 0:\n es = self.g22.encode(fixed_svhn)\n fake_mnist_var = self.g22.decode(es, mnist=True)\n fake_mnist = self.to_data(fake_mnist_var)\n if self.config.save_models_and_samples:\n merged = self.merge_images(svhn_fixed_data, fake_mnist)\n path = os.path.join(self.sample_path, 'sample-%d-s-m.png' % (step + 1))\n scipy.misc.imsave(path, merged)\n print('saved %s' % path)\n\n if (step + 1) % self.config.num_iters_save_model_and_return == 0:\n # save the model parameters for each epoch\n if self.config.save_models_and_samples:\n g22_path = os.path.join(self.model_path, 'g22-%d.pkl' % (step + 1))\n d1_path = os.path.join(self.model_path, 'd1-%d.pkl' % (step + 1))\n d2_path = os.path.join(self.model_path, 'd2-%d.pkl' % (step + 1))\n torch.save(self.g22.state_dict(), g22_path)\n torch.save(self.d1.state_dict(), d1_path)\n torch.save(self.d2.state_dict(), d2_path)\n\n return\n" } ]