{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n \nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n for filename in filenames:\n print(os.path.join(dirname, filename))\n\n# You can write up to 5GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install git+https://github.com/pabloppp/pytorch-tools -U","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"\nimport os\n\n\nimport torch\nimport torch.nn as nn\n\nimport torch.nn.functional as F\nfrom fastai.vision import *\nimport torch.utils.data as data\nfrom torch.utils.data import DataLoader\n\nimport torchvision.transforms as transforms\nfrom PIL import Image as I\nimport torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport os\nimport cv2\nimport torchvision\nimport math\nimport torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.nn import Parameter as P\nfrom functools import partial\nfrom torchtools.optim import RangerLars\nfrom torchvision import datasets\n\nimport math\nfrom torch.optim.optimizer import Optimizer\nimport albumentations as A\nfrom albumentations.pytorch import ToTensor\nfrom numba import jit, njit","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nclass MyImg():\n def __init__(self, img, tfm):\n self.px = np.array(img)\n self.tfm = tfm\n \n @property\n def size(self):\n h, w, _ = self.px.shape\n return min(w, h)\n\n\ndef pad(img, padding_mode='reflect'):\n p = math.ceil((max(img.size) - min(img.size)) / 2)\n p_horr = p if img.width < img.height else 0\n p_vert = p if img.height < img.width else 0\n img = T.Pad((p_horr, p_vert), padding_mode=padding_mode)(img)\n if img.width != img.height:\n s = min(img.size)\n img = img.crop((0, 0, s, s))\n return img\n\n\ndef take_top(img):\n size = min(img.size)\n bbox = (0, 0, size, size)\n return img.crop(bbox)\n\n\ndef take_diagonal(img):\n w, h = img.size\n size = min(w, h)\n bbox_l = (0, 0, size, size)\n bbox_r = (w - size, h - size, w, h)\n return [img.crop(bbox_l), img.crop(bbox_r)]\n \nimport torchvision.transforms as T\n\nIMG_SIZE = 64\nIMG_SIZE_2 = IMG_SIZE * 2\nIMG_CHANNELS = 3\nIMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')\nresize = T.Resize(IMG_SIZE, interpolation=I.LANCZOS)\nresize2x = T.Resize(IMG_SIZE_2, interpolation=I.LANCZOS)\n\ntop_crop = T.Compose([T.Lambda(take_top), resize])\ntop_crop2x = T.Compose([T.Lambda(take_top), resize2x])\n\ntwo_crops = T.Compose([resize, T.Lambda(take_diagonal)])\ntwo_crops2x = T.Compose([resize2x, T.Lambda(take_diagonal)])\n\npad_only = T.Compose([T.Lambda(pad), resize])\npad_only2x = T.Compose([T.Lambda(pad), resize2x])\n\n\n@jit(nopython=True)\ndef pad_one_dim(clow, chigh, pad, cmax):\n clow = max(0, clow - pad)\n chigh = min(cmax, chigh + pad)\n return clow, chigh, chigh - clow\n\n\ndef calc_bbox(obj, img_w, img_h, zoom=0.0, try_square=True):\n bndbox = obj.find('bndbox')\n xmin = int(bndbox.find('xmin').text)\n ymin = int(bndbox.find('ymin').text)\n xmax = int(bndbox.find('xmax').text)\n ymax = int(bndbox.find('ymax').text)\n \n # occasionally i get bboxes which exceed img size\n xmin, xmax, obj_w = pad_one_dim(xmin, xmax, 0, img_w)\n ymin, ymax, obj_h = pad_one_dim(ymin, ymax, 0, img_h)\n \n if zoom != 0.0:\n pad_w = obj_w * zoom / 2\n pad_h = obj_h * zoom / 2\n xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad_w, img_w)\n ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad_h, img_h)\n \n if try_square:\n # try pad both sides equaly\n if obj_w > obj_h:\n pad = (obj_w - obj_h) / 2\n ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad, img_h)\n elif obj_h > obj_w:\n pad = (obj_h - obj_w) / 2\n xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad, img_w)\n\n # if it's still not square, try pad where possible\n if obj_w > obj_h:\n pad = obj_w - obj_h\n ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad, img_h)\n elif obj_h > obj_w:\n pad = obj_h - obj_w\n xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad, img_w)\n \n return int(xmin), int(ymin), int(xmax), int(ymax)\n\n\n@jit(nopython=True)\ndef bb2wh(bbox):\n width = bbox[2] - bbox[0]\n height = bbox[3] - bbox[1]\n return width, height\n\n\ndef create_counters():\n return { 'top': 0, 'sides': 0, 'slide': 0, 'pad': 0, 'skip': 0, 'noop': 0 }\n\n\ndef make_x1res(img, bbox, counter):\n ar = img.width / img.height\n if ar == 1.0:\n tfm_imgs = resize(img); counter['noop'] += 1\n elif 1.0 < ar < 1.4 or 1.0 < 1/ar < 1.25:\n tfm_imgs = pad_only(img); counter['pad'] += 1\n elif 1.4 <= ar < 2.0:\n tfm_imgs = two_crops(img); counter['sides'] += 2\n elif 1.25 <= 1/ar < 1.5:\n tfm_imgs = [pad_only(img), top_crop(img)]\n counter['pad'] += 1; counter['top'] += 1\n elif 1.5 <= 1/ar < 1.8:\n tfm_imgs = top_crop(img); counter['top'] += 1\n else:\n tfm_imgs = None; counter['skip'] += 1\n return tfm_imgs\n\n\ndef make_x2res(img, bbox, counter):\n ar = img.width / img.height\n if ar == 1.0: \n tfm_imgs = resize2x(img); counter['noop'] += 1\n elif 1.0 < ar < 1.4 or 1.0 < 1/ar < 1.4:\n tfm_imgs = pad_only2x(img); counter['pad'] += 1\n else:\n tfm_imgs = None\n return tfm_imgs\n \n \ndef add_sample(samples, label, tfm, imgs, labels):\n if not samples: return\n elif isinstance(samples, I.Image):\n imgs.append(MyImg(samples, tfm))\n labels.append(label)\n elif isinstance(samples, list):\n imgs.extend([MyImg(s, tfm) for s in samples])\n labels.extend([label] * len(samples))\n else: assert False\n\n \ndef is_valid_file(x):\n return datasets.folder.has_file_allowed_extension(x, IMG_EXTENSIONS)\n \n\nclass DogsDataSet(datasets.vision.VisionDataset):\n def __init__(self, root, transforms, target_transform=None, max_samples=None):\n super().__init__(root, transform=None)\n assert isinstance(transforms, list) and len(transforms) == 4\n self.transforms = transforms\n self.target_transform = target_transform\n self.max_samples = max_samples\n self.classes = {}\n \n imgs, labels = self._load_subfolders_images(self.root) \n assert len(imgs) == len(labels)\n if len(imgs) == 0:\n raise RuntimeError(f'Found 0 files in subfolders of: {self.root}')\n self.imgs = imgs\n self.labels = labels\n \n def _create_or_get_class(self, name):\n try:\n label = self.classes[name]\n except KeyError:\n label = len(self.classes)\n self.classes[name] = label\n return label\n \n def _load_subfolders_images(self, root):\n imgs, labels, paths = [], [], []\n add_sample_ = partial(add_sample, imgs=imgs, labels=labels)\n low_res = lambda bbox: min(bb2wh(bbox)) < IMG_SIZE\n can_2x = lambda bbox: min(bb2wh(bbox)) >= IMG_SIZE_2\n counter = create_counters()\n\n for root, _, fnames in sorted(os.walk(root)):\n for fname in sorted(fnames):\n path = os.path.join(root, fname)\n paths.append(path)\n if self.max_samples:\n paths = paths[:self.max_samples]\n\n for path in paths:\n if not is_valid_file(path): continue\n img = datasets.folder.default_loader(path)\n annotation_basename = os.path.splitext(os.path.basename(path))[0]\n annotation_dirname = next(dirname for dirname in os.listdir(ANNOT_PATH) if dirname.startswith(annotation_basename.split('_')[0]))\n annotation_filename = os.path.join(ANNOT_PATH, annotation_dirname, annotation_basename)\n tree = ET.parse(annotation_filename)\n root = tree.getroot()\n objects = root.findall('object')\n for o in objects:\n name = o.find('name').text\n label = self._create_or_get_class(name)\n prev_bbox, tfm_imgs = None, None\n\n bbox = calc_bbox(o, img_w=img.width, img_h=img.height, zoom=0.03125) \n if low_res(bbox):\n counter['skip'] += 1; continue\n obj_img = img.crop(bbox)\n ar = obj_img.width / obj_img.height\n tfm_imgs = make_x1res(obj_img, bbox, counter)\n add_sample_(tfm_imgs, label, 0)\n# if can_2x(bbox):\n# tfm_imgs = make_x2res(obj_img, bbox, counter)\n# add_sample_(tfm_imgs, label, 1)\n# prev_bbox = bbox\n \n# bbox = calc_bbox(o, img_w=img.width, img_h=img.height, zoom=0.1)\n# if not (bbox == prev_bbox or low_res(bbox)):\n# obj_img = img.crop(bbox)\n# ar = obj_img.width / obj_img.height\n# if 1.05 < ar < 1.25:\n# tfm_imgs = resize(obj_img); counter['slide'] += 1\n# add_sample_(tfm_imgs, label, 3)\n# if can_2x(bbox):\n# tfm_imgs = make_x2res(obj_img, bbox, counter)\n# add_sample_(tfm_imgs, label, 2)\n\n# bbox = calc_bbox(o, img_w=img.width, img_h=img.height, zoom=0.12, try_square=False)\n# if not low_res(bbox):\n# obj_img = img.crop(bbox)\n# ar = obj_img.width / obj_img.height\n# if 1.3 <= 1/ar < 1.8:\n# tfm_imgs = top_crop(obj_img); counter['top'] += 1\n# add_sample_(tfm_imgs, label, 0)\n# if can_2x(bbox):\n# tfm_imgs = make_x2res(obj_img, bbox, counter)\n# add_sample_(tfm_imgs, label, 1)\n \n n_x1, n_x2 = 0, 0\n for i, img in enumerate(imgs):\n if img.size == IMG_SIZE: \n n_x1 +=1\n else: \n n_x2 +=1\n \n print(f'Found {len(self.classes)} classes\\nLoaded 64x64 {n_x1} images\\n'\n f'Loaded 128x128 {n_x2} images\\n')\n print(counter)\n return imgs, labels\n \n def __getitem__(self, index):\n img = self.imgs[index]\n label = self.labels[index]\n tfms = self.transforms[img.tfm]\n img = tfms(image=img.px)['image'] \n if self.target_transform:\n label = self.target_transform(label)\n return img, label\n\n def __len__(self):\n return len(self.imgs)\n \n \ndef create_runtime_tfms():\n mean, std = [0.5]*3, [0.5]*3\n resize_to_64 = A.SmallestMaxSize(IMG_SIZE, interpolation=cv2.INTER_LANCZOS4)\n out = [A.HorizontalFlip(p=0.5), A.Normalize(mean=mean, std=std), ToTensor()]\n rand_crop = A.RandomCrop(IMG_SIZE, IMG_SIZE)\n\n affine_1 = A.ShiftScaleRotate(\n shift_limit=0, scale_limit=0.1, rotate_limit=8, \n interpolation=cv2.INTER_CUBIC, \n border_mode=cv2.BORDER_REFLECT_101, p=1.0)\n affine_1 = A.Compose([affine_1, resize_to_64])\n\n affine_2 = A.ShiftScaleRotate(\n shift_limit=0.06, scale_limit=(-0.06, 0.14), rotate_limit=6, \n interpolation=cv2.INTER_CUBIC, \n border_mode=cv2.BORDER_REFLECT_101, p=1.0)\n affine_2 = A.Compose([affine_2, resize_to_64])\n\n tfm_0 = A.Compose(out)\n tfm_1 = A.Compose([affine_1, *out])\n tfm_2 = A.Compose([affine_2, *out])\n tfm_3 = A.Compose([rand_crop, *out])\n return [tfm_0, tfm_1, tfm_2, tfm_3]\n\n\ndef get_data_loaders(data_root=None, batch_size=32, num_workers=2, shuffle=True,\n pin_memory=True, drop_last=True):\n print('Using dataset root location %s' % data_root)\n # tfms_x1, tfms_x2 = create_runtime_tfms()\n train_set = DogsDataSet(data_root, create_runtime_tfms())\n # Prepare loader; the loaders list is for forward compatibility with\n # using validation / test splits.\n loaders = []\n loader_kwargs = {'num_workers': num_workers, 'pin_memory': pin_memory,\n 'drop_last': drop_last} # Default, drop last incomplete batch\n train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=shuffle, **loader_kwargs)\n loaders.append(train_loader)\n return loaders","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"IMAGE_SIZE=64\ndevice=torch.device('cuda:0')\ndef denorm(x):\n out = (x + 1) / 2\n return out.clamp_(0, 1)\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport cv2\nimport xml.etree.ElementTree as ET\n\n\ndef pad_one_dim(clow, chigh, pad, cmax):\n clow = max(0, clow - pad)\n chigh = min(cmax, chigh + pad)\n return clow, chigh, chigh - clow\n\n\ndef calc_bbox(obj, img_w, img_h, zoom=0.0, try_square=True):\n bndbox = obj.find('bndbox')\n xmin = int(bndbox.find('xmin').text)\n ymin = int(bndbox.find('ymin').text)\n xmax = int(bndbox.find('xmax').text)\n\n ymax = int(bndbox.find('ymax').text)\n\n # occasionally i get bboxes which exceed img size\n xmin, xmax, obj_w = pad_one_dim(xmin, xmax, 0, img_w)\n ymin, ymax, obj_h = pad_one_dim(ymin, ymax, 0, img_h)\n\n if zoom != 0.0:\n pad_w = obj_w * zoom / 2\n pad_h = obj_h * zoom / 2\n xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad_w, img_w)\n ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad_h, img_h)\n\n if try_square:\n # try pad both sides equaly\n if obj_w > obj_h:\n pad = (obj_w - obj_h) / 2\n ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad, img_h)\n elif obj_h > obj_w:\n pad = (obj_h - obj_w) / 2\n xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad, img_w)\n\n # if it's still not square, try pad where possible\n if obj_w > obj_h:\n pad = obj_w - obj_h\n ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad, img_h)\n elif obj_h > obj_w:\n pad = obj_h - obj_w\n xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad, img_w)\n\n return int(xmin), int(ymin), int(xmax), int(ymax)\n\n\n\n\ndef make_bbox(image, bbox):\n x0, y0, x1, y1 = bbox\n image = image[y0:y1, x0:x1, :]\n return image\n\n\nDog_transform=A.Compose([\n A.Resize(IMAGE_SIZE,IMAGE_SIZE),\n A.HorizontalFlip(),\n A.Rotate(45),\n A.Normalize(),\n ToTensorV2()\n])\nclass DogDataset(data.Dataset):\n def __init__(self,image_paths,anno_paths,transform):\n super(DogDataset, self).__init__()\n self.image_paths=image_paths\n self.anno_paths=anno_paths\n self.transform=transform\n self.image_labels={}\n\n def make_label(self,name):\n try:\n label = self.image_labels[name]\n except KeyError:\n label = len(self.image_labels)\n self.image_labels[name] = label\n return label\n\n def __getitem__(self,index):\n path=self.image_paths[index]\n img = cv2.imread(str(path))\n \n\n annotation_basename = os.path.splitext(os.path.basename(path))[0]\n annotation_dirname = next(dirname for dirname in os.listdir(self.anno_paths) if\n dirname.startswith(annotation_basename.split('_')[0]))\n annotation_filename = os.path.join(self.anno_paths, annotation_dirname, annotation_basename)\n tree = ET.parse(annotation_filename)\n root = tree.getroot()\n objects = root.findall('object')\n\n o=objects[0]\n name = o.find('name').text\n label = self.make_label(name)\n\n bbox = calc_bbox(o, img_w=img.shape[1], img_h=img.shape[0], zoom=0.03125)\n image = make_bbox(img, bbox)\n if image.shape[0]==0:\n print(bbox)\n print(img.shape)\n auged=self.transform(image=image)\n image=auged['image']\n \n return image,label\n \n def __len__(self):\n return len(self.image_paths)\n \n \ndef get_dogloader(image_paths,anno_paths,batch_size=32,shuffle=True,drop_last=True):\n dataset=DogDataset(image_paths,anno_paths,Dog_transform)\n print(len(dataset))\n dataloader=data.DataLoader(dataset,batch_size=batch_size,shuffle=shuffle,drop_last=drop_last)\n return dataloader\n\n\nclass myDataset(data.Dataset):\n def __init__(self,train_names,image_labels,train_transform,bbox=None):\n super(myDataset, self).__init__()\n self.train_names=train_names\n self.image_labels=image_labels\n self.train_transform=train_transform\n self.bbox=bbox\n\n def __getitem__(self,index):\n name=self.train_names[index]\n label=self.image_labels[name]\n\n image=I.open(name)\n image=image.convert('RGB')\n if self.bbox is not None:\n bbox=self.bbox[name]\n x0,x1,y0,y1=bbox\n image=image[x0:x1,y0:y1,:]\n \n image=self.train_transform(image)\n return image,label\n \n def __len__(self):\n return len(self.train_names)\n\ndef get_transform(Image_size):\n return transforms.Compose([\n transforms.Resize([Image_size,Image_size]),\n# transforms.ColorJitter(),\n transforms.RandomHorizontalFlip(),\n transforms.RandomRotation(45),\n transforms.ToTensor(),\n transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n# transforms.RandomErasing()\n ])\n\n\n\n\n\n\n\n\n\ndef get_dataloader(train_names,image_labels,batch_size=32,shuffle=True,drop_last=True,bbox=None):\n transform=get_transform(IMAGE_SIZE)\n dataset=myDataset(train_names,image_labels,transform,bbox)\n dataloader=data.DataLoader(dataset,batch_size=batch_size,shuffle=shuffle,drop_last=drop_last)\n return dataloader\n\ndef seed_rng(seed):\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n np.random.seed(seed)\n\n\ndef plot_imgs(imgs, cols=8, size=4):\n n = imgs.shape[0]\n rows = n // cols\n _, axes = plt.subplots(figsize=(cols * size, rows * size), ncols=cols, nrows=rows)\n for i, ax in enumerate(axes.flatten()):\n img = denorm(imgs[i]).numpy().transpose(1, 2, 0)\n ax.get_xaxis().set_visible(False)\n ax.get_yaxis().set_visible(False)\n ax.imshow(img)\n plt.subplots_adjust(wspace=0.0, hspace=0.0)\n plt.show()\n\ndef save_and_sample(G, G_ema, fixed_z, fixed_y, config):\n which_G = G_ema if config['ema'] and config['use_ema'] else G\n with torch.no_grad():\n fixed_Gz = which_G(fixed_z, which_G.shared(fixed_y))\n sample_images = fixed_Gz.float().cpu()\n plot_imgs(sample_images, size=2)\n\n\n\n'''限制噪声'''\ndef trunc_trick(bs, z_dim, bound=2.0):\n z = torch.randn(bs, z_dim)\n while z.min() < -bound or bound < z.max():\n z = z.where((-bound < z) & (z < bound), torch.randn_like(z))\n return z\n\ndef collect_bn_stats(G, n_samples, config):\n im_batch_size = config['n_classes']\n G.train()\n\n for i_batch in range(0, n_samples, im_batch_size):\n with torch.no_grad():\n z = torch.randn(im_batch_size, G.dim_z, device=device)\n y = torch.arange(im_batch_size).to(device)\n images = G(z, G.shared(y)).float().cpu()\n\n\n'''生成图片'''\ndef generate_images(out_dir, G, n_images, config):\n im_batch_size = config['n_classes']\n z_bound = config['trunc_z']\n if z_bound > 0.0: print(f'Truncating z to (-{z_bound}, {z_bound})')\n\n for i_batch in range(0, n_images, im_batch_size):\n with torch.no_grad():\n if z_bound > 0.0:\n z = trunc_trick(im_batch_size, G.dim_z, bound=z_bound).to(device)\n else:\n z = torch.randn(im_batch_size, G.dim_z, device=device)\n y = torch.arange(im_batch_size).to(device)\n images = G(z, G.shared(y)).float().cpu()\n\n if i_batch + im_batch_size > n_images:\n n_last_images = n_images - i_batch\n print(f'Taking only {n_last_images} images from the last batch...')\n images = images[:n_last_images]\n\n for i_image, image in enumerate(images):\n fname = os.path.join(out_dir, f'image_{i_batch+i_image:05d}.png')\n image = denorm(image)\n if config['denoise']:\n image = image * 255\n image = image.numpy().transpose(1, 2, 0).astype('uint8')\n image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)\n image = cv2.fastNlMeansDenoisingColored(\n src=image, dst=None,\n h=config['denoise_str_lum'],\n hColor=config['denoise_str_col'],\n templateWindowSize=config['denoise_kernel_size'],\n searchWindowSize=config['denoise_search_window']\n )\n cv2.imwrite(fname, image)\n else:\n torchvision.utils.save_image(image, fname)\n\ndef show_saved_samples(output_dir,n_images=256, cols=8, size=3):\n rows = math.ceil(n_images / cols)\n _, axes = plt.subplots(rows, cols, figsize=(cols * size, rows * size))\n for i, ax in enumerate(axes.flatten()):\n img = I.open(f'{output_dir}/image_{i:05d}.png')\n ax.get_xaxis().set_visible(False)\n ax.get_yaxis().set_visible(False)\n ax.imshow(img)\n plt.subplots_adjust(wspace=0.0, hspace=0.0)\n plt.show()\n\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nclass Adam16(Optimizer):\n def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):\n defaults = dict(lr=lr, betas=betas, eps=eps,\n weight_decay=weight_decay)\n params = list(params)\n super(Adam16, self).__init__(params, defaults)\n\n # Safety modification to make sure we floatify our state\n def load_state_dict(self, state_dict):\n super(Adam16, self).load_state_dict(state_dict)\n for group in self.param_groups:\n for p in group['params']:\n self.state[p]['exp_avg'] = self.state[p]['exp_avg'].float()\n self.state[p]['exp_avg_sq'] = self.state[p]['exp_avg_sq'].float()\n self.state[p]['fp32_p'] = self.state[p]['fp32_p'].float()\n\n def step(self, closure=None):\n \"\"\"Performs a single optimization step.\n Arguments:\n closure (callable, optional): A closure that reevaluates the model\n and returns the loss.\n \"\"\"\n loss = None\n if closure is not None:\n loss = closure()\n\n for group in self.param_groups:\n for p in group['params']:\n if p.grad is None:\n continue\n\n grad = p.grad.data.float()\n state = self.state[p]\n\n # State initialization\n if len(state) == 0:\n state['step'] = 0\n # Exponential moving average of gradient values\n state['exp_avg'] = grad.new().resize_as_(grad).zero_()\n # Exponential moving average of squared gradient values\n state['exp_avg_sq'] = grad.new().resize_as_(grad).zero_()\n # Fp32 copy of the weights\n state['fp32_p'] = p.data.float()\n\n exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n beta1, beta2 = group['betas']\n\n state['step'] += 1\n\n if group['weight_decay'] != 0:\n grad = grad.add(group['weight_decay'], state['fp32_p'])\n\n # Decay the first and second moment running average coefficient\n exp_avg.mul_(beta1).add_(1 - beta1, grad)\n exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n\n denom = exp_avg_sq.sqrt().add_(group['eps'])\n\n bias_correction1 = 1 - beta1 ** state['step']\n bias_correction2 = 1 - beta2 ** state['step']\n step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1\n\n state['fp32_p'].addcdiv_(-step_size, exp_avg, denom)\n p.data = state['fp32_p'].half()\n\n return loss\ndef proj(x, y):\n return torch.mm(y, x.t()) * y / torch.mm(y, y.t())\n\n\n# Orthogonalize x wrt list of vectors ys\ndef gram_schmidt(x, ys):\n for y in ys:\n x = x - proj(x, y)\n return x\n\n\n# Apply num_itrs steps of the power method to estimate top N singular values.\ndef power_iteration(W, u_, update=True, eps=1e-12):\n # Lists holding singular vectors and values\n us, vs, svs = [], [], []\n for i, u in enumerate(u_):\n # Run one step of the power iteration\n with torch.no_grad():\n v = torch.matmul(u, W)\n # Run Gram-Schmidt to subtract components of all other singular vectors\n v = F.normalize(gram_schmidt(v, vs), eps=eps)\n # Add to the list\n vs += [v]\n # Update the other singular vector\n u = torch.matmul(v, W.t())\n # Run Gram-Schmidt to subtract components of all other singular vectors\n u = F.normalize(gram_schmidt(u, us), eps=eps)\n # Add to the list\n us += [u]\n if update:\n u_[i][:] = u\n # Compute this singular value and add it to the list\n svs += [torch.squeeze(torch.matmul(torch.matmul(v, W.t()), u.t()))]\n # svs += [torch.sum(F.linear(u, W.transpose(0, 1)) * v)]\n return svs, us, vs\n\n\n# Convenience passthrough function\nclass identity(nn.Module):\n def forward(self, input):\n return input\n\n\n# Spectral normalization base class\nclass SN(object):\n def __init__(self, num_svs, num_itrs, num_outputs, transpose=False, eps=1e-12):\n # Number of power iterations per step\n self.num_itrs = num_itrs\n # Number of singular values\n self.num_svs = num_svs\n # Transposed?\n self.transpose = transpose\n # Epsilon value for avoiding divide-by-0\n self.eps = eps\n # Register a singular vector for each sv\n for i in range(self.num_svs):\n self.register_buffer('u%d' % i, torch.randn(1, num_outputs))\n self.register_buffer('sv%d' % i, torch.ones(1))\n\n # Singular vectors (u side)\n @property\n def u(self):\n return [getattr(self, 'u%d' % i) for i in range(self.num_svs)]\n\n # Singular values;\n # note that these buffers are just for logging and are not used in training.\n @property\n def sv(self):\n return [getattr(self, 'sv%d' % i) for i in range(self.num_svs)]\n\n # Compute the spectrally-normalized weight\n def W_(self):\n W_mat = self.weight.view(self.weight.size(0), -1)\n if self.transpose:\n W_mat = W_mat.t()\n # Apply num_itrs power iterations\n for _ in range(self.num_itrs):\n svs, us, vs = power_iteration(W_mat, self.u, update=self.training, eps=self.eps)\n # Update the svs\n if self.training:\n with torch.no_grad(): # Make sure to do this in a no_grad() context or you'll get memory leaks!\n for i, sv in enumerate(svs):\n self.sv[i][:] = sv\n return self.weight / svs[0]\n\n\n# 2D Conv layer with spectral norm\nclass SNConv2d(nn.Conv2d, SN):\n def __init__(self, in_channels, out_channels, kernel_size, stride=1,\n padding=0, dilation=1, groups=1, bias=True,\n num_svs=1, num_itrs=1, eps=1e-12):\n nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, stride,\n padding, dilation, groups, bias)\n SN.__init__(self, num_svs, num_itrs, out_channels, eps=eps)\n\n def forward(self, x):\n return F.conv2d(x, self.W_(), self.bias, self.stride,\n self.padding, self.dilation, self.groups)\n\n\n# Linear layer with spectral norm\nclass SNLinear(nn.Linear, SN):\n def __init__(self, in_features, out_features, bias=True,\n num_svs=1, num_itrs=1, eps=1e-12):\n nn.Linear.__init__(self, in_features, out_features, bias)\n SN.__init__(self, num_svs, num_itrs, out_features, eps=eps)\n\n def forward(self, x):\n return F.linear(x, self.W_(), self.bias)\n\n\n# Embedding layer with spectral norm\n# We use num_embeddings as the dim instead of embedding_dim here\n# for convenience sake\nclass SNEmbedding(nn.Embedding, SN):\n def __init__(self, num_embeddings, embedding_dim, padding_idx=None,\n max_norm=None, norm_type=2, scale_grad_by_freq=False,\n sparse=False, _weight=None,\n num_svs=1, num_itrs=1, eps=1e-12):\n nn.Embedding.__init__(self, num_embeddings, embedding_dim, padding_idx,\n max_norm, norm_type, scale_grad_by_freq,\n sparse, _weight)\n SN.__init__(self, num_svs, num_itrs, num_embeddings, eps=eps)\n\n def forward(self, x):\n return F.embedding(x, self.W_())\n\n\n# A non-local block as used in SA-GAN\n# Note that the implementation as described in the paper is largely incorrect;\n# refer to the released code for the actual implementation.\n# class Attention(nn.Module):\n# def __init__(self, ch, which_conv=SNConv2d, name='attention'):\n# super(Attention, self).__init__()\n# # Channel multiplier\n# self.ch = ch\n# self.which_conv = which_conv\n# self.theta = self.which_conv(self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False)\n# self.phi = self.which_conv(self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False)\n# self.g = self.which_conv(self.ch, self.ch // 2, kernel_size=1, padding=0, bias=False)\n# self.o = self.which_conv(self.ch // 2, self.ch, kernel_size=1, padding=0, bias=False)\n# # Learnable gain parameter\n# self.gamma = P(torch.tensor(0.), requires_grad=True)\n\n# def forward(self, x, y=None):\n# # Apply convs\n# theta = self.theta(x)\n# phi = F.max_pool2d(self.phi(x), [2, 2])\n# g = F.max_pool2d(self.g(x), [2, 2])\n# # Perform reshapes\n# theta = theta.view(-1, self.ch // 8, x.shape[2] * x.shape[3])\n# phi = phi.view(-1, self.ch // 8, x.shape[2] * x.shape[3] // 4)\n# g = g.view(-1, self.ch // 2, x.shape[2] * x.shape[3] // 4)\n# # Matmul and softmax to get attention maps\n# beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), -1)\n# # Attention map times g path\n# o = self.o(torch.bmm(g, beta.transpose(1, 2)).view(-1, self.ch // 2, x.shape[2], x.shape[3]))\n# return self.gamma * o + x\n\n\nclass Attention(nn.Module):\n def __init__(self, ch, which_conv=SNConv2d, name='attention'):\n super(Attention, self).__init__()\n # Channel multiplier\n self.ch = ch\n self.which_conv = which_conv\n self.theta = self.which_conv(self.ch, self.ch // 4, kernel_size=1, padding=0, bias=False)\n self.phi = self.which_conv(self.ch, self.ch // 4, kernel_size=1, padding=0, bias=False)\n self.g = self.which_conv(self.ch, self.ch // 2, kernel_size=1, padding=0, bias=False)\n self.o = self.which_conv(self.ch // 2, self.ch, kernel_size=1, padding=0, bias=False)\n # Learnable gain parameter\n self.gamma = P(torch.tensor(0.), requires_grad=True)\n\n def forward(self, x, y=None):\n # Apply convs\n theta = self.theta(x)\n phi = F.max_pool2d(self.phi(x), [2, 2])\n g = F.max_pool2d(self.g(x), [2, 2])\n # Perform reshapes\n theta = theta.view(-1, self.ch // 4, x.shape[2] * x.shape[3])\n phi = phi.view(-1, self.ch // 4, x.shape[2] * x.shape[3] // 4)\n g = g.view(-1, self.ch // 2, x.shape[2] * x.shape[3] // 4)\n # Matmul and softmax to get attention maps\n beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), -1)\n # Attention map times g path\n o = self.o(torch.bmm(g, beta.transpose(1, 2)).view(-1, self.ch // 2, x.shape[2], x.shape[3]))\n return self.gamma * o + x\n\n\n# Fused batchnorm op\ndef fused_bn(x, mean, var, gain=None, bias=None, eps=1e-5):\n # Apply scale and shift--if gain and bias are provided, fuse them here\n # Prepare scale\n scale = torch.rsqrt(var + eps)\n # If a gain is provided, use it\n if gain is not None:\n scale = scale * gain\n # Prepare shift\n shift = mean * scale\n # If bias is provided, use it\n if bias is not None:\n shift = shift - bias\n return x * scale - shift\n # return ((x - mean) / ((var + eps) ** 0.5)) * gain + bias # The unfused way.\n\n\n# Manual BN\n# Calculate means and variances using mean-of-squares minus mean-squared\ndef manual_bn(x, gain=None, bias=None, return_mean_var=False, eps=1e-5):\n # Cast x to float32 if necessary\n float_x = x.float()\n # Calculate expected value of x (m) and expected value of x**2 (m2)\n # Mean of x\n m = torch.mean(float_x, [0, 2, 3], keepdim=True)\n # Mean of x squared\n m2 = torch.mean(float_x ** 2, [0, 2, 3], keepdim=True)\n # Calculate variance as mean of squared minus mean squared.\n var = (m2 - m ** 2)\n # Cast back to float 16 if necessary\n var = var.type(x.type())\n m = m.type(x.type())\n # Return mean and variance for updating stored mean/var if requested\n if return_mean_var:\n return fused_bn(x, m, var, gain, bias, eps), m.squeeze(), var.squeeze()\n else:\n return fused_bn(x, m, var, gain, bias, eps)\n\n\n# My batchnorm, supports standing stats\nclass myBN(nn.Module):\n def __init__(self, num_channels, eps=1e-5, momentum=0.1):\n super(myBN, self).__init__()\n # momentum for updating running stats\n self.momentum = momentum\n # epsilon to avoid dividing by 0\n self.eps = eps\n # Momentum\n self.momentum = momentum\n # Register buffers\n self.register_buffer('stored_mean', torch.zeros(num_channels))\n self.register_buffer('stored_var', torch.ones(num_channels))\n self.register_buffer('accumulation_counter', torch.zeros(1))\n # Accumulate running means and vars\n self.accumulate_standing = False\n\n # reset standing stats\n def reset_stats(self):\n self.stored_mean[:] = 0\n self.stored_var[:] = 0\n self.accumulation_counter[:] = 0\n\n def forward(self, x, gain, bias):\n if self.training:\n out, mean, var = manual_bn(x, gain, bias, return_mean_var=True, eps=self.eps)\n # If accumulating standing stats, increment them\n if self.accumulate_standing:\n self.stored_mean[:] = self.stored_mean + mean.data\n self.stored_var[:] = self.stored_var + var.data\n self.accumulation_counter += 1.0\n # If not accumulating standing stats, take running averages\n else:\n self.stored_mean[:] = self.stored_mean * (1 - self.momentum) + mean * self.momentum\n self.stored_var[:] = self.stored_var * (1 - self.momentum) + var * self.momentum\n return out\n # If not in training mode, use the stored statistics\n else:\n mean = self.stored_mean.view(1, -1, 1, 1)\n var = self.stored_var.view(1, -1, 1, 1)\n # If using standing stats, divide them by the accumulation counter\n if self.accumulate_standing:\n mean = mean / self.accumulation_counter\n var = var / self.accumulation_counter\n return fused_bn(x, mean, var, gain, bias, self.eps)\n\n\n# Simple function to handle groupnorm norm stylization\ndef groupnorm(x, norm_style):\n # If number of channels specified in norm_style:\n if 'ch' in norm_style:\n ch = int(norm_style.split('_')[-1])\n groups = max(int(x.shape[1]) // ch, 1)\n # If number of groups specified in norm style\n elif 'grp' in norm_style:\n groups = int(norm_style.split('_')[-1])\n # If neither, default to groups = 16\n else:\n groups = 16\n return F.group_norm(x, groups)\n\n\n# Class-conditional bn\n# output size is the number of channels, input size is for the linear layers\n# Andy's Note: this class feels messy but I'm not really sure how to clean it up\n# Suggestions welcome! (By which I mean, refactor this and make a pull request\n# if you want to make this more readable/usable).\nclass ccbn(nn.Module):\n def __init__(self, output_size, input_size, which_linear, eps=1e-5, momentum=0.1,\n cross_replica=False, mybn=False, norm_style='bn', ):\n super(ccbn, self).__init__()\n self.output_size, self.input_size = output_size, input_size\n # Prepare gain and bias layers\n self.gain = which_linear(input_size, output_size)\n self.bias = which_linear(input_size, output_size)\n # epsilon to avoid dividing by 0\n self.eps = eps\n # Momentum\n self.momentum = momentum\n # Use cross-replica batchnorm?\n self.cross_replica = cross_replica\n # Use my batchnorm?\n self.mybn = mybn\n # Norm style?\n self.norm_style = norm_style\n\n if self.mybn:\n self.bn = myBN(output_size, self.eps, self.momentum)\n elif self.norm_style in ['bn', 'in']:\n self.register_buffer('stored_mean', torch.zeros(output_size))\n self.register_buffer('stored_var', torch.ones(output_size))\n\n def forward(self, x, y):\n # Calculate class-conditional gains and biases\n gain = (1 + self.gain(y)).view(y.size(0), -1, 1, 1)\n bias = self.bias(y).view(y.size(0), -1, 1, 1)\n # If using my batchnorm\n if self.mybn:\n return self.bn(x, gain=gain, bias=bias)\n # else:\n else:\n if self.norm_style == 'bn':\n out = F.batch_norm(x, self.stored_mean, self.stored_var, None, None,\n self.training, 0.1, self.eps)\n elif self.norm_style == 'in':\n out = F.instance_norm(x, self.stored_mean, self.stored_var, None, None,\n self.training, 0.1, self.eps)\n elif self.norm_style == 'gn':\n out = groupnorm(x, self.normstyle)\n elif self.norm_style == 'nonorm':\n out = x\n return out * gain + bias\n\n def extra_repr(self):\n s = 'out: {output_size}, in: {input_size},'\n s += ' cross_replica={cross_replica}'\n return s.format(**self.__dict__)\n\n\n# Normal, non-class-conditional BN\nclass bn(nn.Module):\n def __init__(self, output_size, eps=1e-5, momentum=0.1,\n cross_replica=False, mybn=False):\n super(bn, self).__init__()\n self.output_size = output_size\n # Prepare gain and bias layers\n self.gain = P(torch.ones(output_size), requires_grad=True)\n self.bias = P(torch.zeros(output_size), requires_grad=True)\n # epsilon to avoid dividing by 0\n self.eps = eps\n # Momentum\n self.momentum = momentum\n # Use cross-replica batchnorm?\n self.cross_replica = cross_replica\n # Use my batchnorm?\n self.mybn = mybn\n\n if mybn:\n self.bn = myBN(output_size, self.eps, self.momentum)\n # Register buffers if neither of the above\n else:\n self.register_buffer('stored_mean', torch.zeros(output_size))\n self.register_buffer('stored_var', torch.ones(output_size))\n\n def forward(self, x, y=None):\n if self.mybn:\n gain = self.gain.view(1, -1, 1, 1)\n bias = self.bias.view(1, -1, 1, 1)\n return self.bn(x, gain=gain, bias=bias)\n else:\n return F.batch_norm(x, self.stored_mean, self.stored_var, self.gain,\n self.bias, self.training, self.momentum, self.eps)\n\n\n# Generator blocks\n# Note that this class assumes the kernel size and padding (and any other\n# settings) have been selected in the main generator module and passed in\n# through the which_conv arg. Similar rules apply with which_bn (the input\n# size [which is actually the number of channels of the conditional info] must\n# be preselected)\nclass GBlock(nn.Module):\n def __init__(self, in_channels, out_channels,\n which_conv=nn.Conv2d, which_bn=bn, activation=None,\n upsample=None):\n super(GBlock, self).__init__()\n\n self.in_channels, self.out_channels = in_channels, out_channels\n self.which_conv, self.which_bn = which_conv, which_bn\n self.activation = activation\n self.upsample = upsample\n # Conv layers\n self.conv1 = self.which_conv(self.in_channels, self.out_channels)\n self.conv2 = self.which_conv(self.out_channels, self.out_channels)\n self.learnable_sc = in_channels != out_channels or upsample\n if self.learnable_sc:\n self.conv_sc = self.which_conv(in_channels, out_channels,\n kernel_size=1, padding=0)\n # Batchnorm layers\n self.bn1 = self.which_bn(in_channels)\n self.bn2 = self.which_bn(out_channels)\n # upsample layers\n self.upsample = upsample\n\n def forward(self, x, y):\n h = self.activation(self.bn1(x, y))\n if self.upsample:\n h = self.upsample(h)\n x = self.upsample(x)\n h = self.conv1(h)\n h = self.activation(self.bn2(h, y))\n h = self.conv2(h)\n if self.learnable_sc:\n x = self.conv_sc(x)\n return h + x\n\n\n# Residual block for the discriminator\nclass DBlock(nn.Module):\n def __init__(self, in_channels, out_channels, which_conv=SNConv2d, wide=True,\n preactivation=False, activation=None, downsample=None, ):\n super(DBlock, self).__init__()\n self.in_channels, self.out_channels = in_channels, out_channels\n # If using wide D (as in SA-GAN and BigGAN), change the channel pattern\n self.hidden_channels = self.out_channels if wide else self.in_channels\n self.which_conv = which_conv\n self.preactivation = preactivation\n self.activation = activation\n self.downsample = downsample\n\n # Conv layers\n self.conv1 = self.which_conv(self.in_channels, self.hidden_channels)\n self.conv2 = self.which_conv(self.hidden_channels, self.out_channels)\n self.learnable_sc = True if (in_channels != out_channels) or downsample else False\n if self.learnable_sc:\n self.conv_sc = self.which_conv(in_channels, out_channels,\n kernel_size=1, padding=0)\n\n def shortcut(self, x):\n if self.preactivation:\n if self.learnable_sc:\n x = self.conv_sc(x)\n if self.downsample:\n x = self.downsample(x)\n else:\n if self.downsample:\n x = self.downsample(x)\n if self.learnable_sc:\n x = self.conv_sc(x)\n return x\n\n def forward(self, x):\n if self.preactivation:\n # h = self.activation(x) # NOT TODAY SATAN\n # Andy's note: This line *must* be an out-of-place ReLU or it\n # will negatively affect the shortcut connection.\n h = F.relu(x)\n else:\n h = x\n h = self.conv1(h)\n h = self.conv2(self.activation(h))\n if self.downsample:\n h = self.downsample(h)\n\n return h + self.shortcut(x)\n\n\ndef G_arch(ch=64, attention='64', ksize='333333', dilation='111111'):\n arch = {}\n arch[512] = {'in_channels': [ch * item for item in [16, 16, 8, 8, 4, 2, 1]],\n 'out_channels': [ch * item for item in [16, 8, 8, 4, 2, 1, 1]],\n 'upsample': [True] * 7,\n 'resolution': [8, 16, 32, 64, 128, 256, 512],\n 'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])\n for i in range(3, 10)}}\n arch[256] = {'in_channels': [ch * item for item in [16, 16, 8, 8, 4, 2]],\n 'out_channels': [ch * item for item in [16, 8, 8, 4, 2, 1]],\n 'upsample': [True] * 6,\n 'resolution': [8, 16, 32, 64, 128, 256],\n 'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])\n for i in range(3, 9)}}\n arch[128] = {'in_channels': [ch * item for item in [16, 16, 8, 4, 2]],\n 'out_channels': [ch * item for item in [16, 8, 4, 2, 1]],\n 'upsample': [True] * 5,\n 'resolution': [8, 16, 32, 64, 128],\n 'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])\n for i in range(3, 8)}}\n arch[64] = {'in_channels': [ch * item for item in [16, 16, 8, 4]],\n 'out_channels': [ch * item for item in [16, 8, 4, 2]],\n 'upsample': [True] * 4,\n 'resolution': [8, 16, 32, 64],\n 'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])\n for i in range(3, 7)}}\n arch[32] = {'in_channels': [ch * item for item in [4, 4, 4]],\n 'out_channels': [ch * item for item in [4, 4, 4]],\n 'upsample': [True] * 3,\n 'resolution': [8, 16, 32],\n 'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])\n for i in range(3, 6)}}\n\n return arch\n\n\nclass Generator(nn.Module):\n def __init__(self, G_ch=64, dim_z=128, bottom_width=4, resolution=128,\n G_kernel_size=3, G_attn='64', n_classes=1000,\n num_G_SVs=1, num_G_SV_itrs=1,\n G_shared=True, shared_dim=0, hier=False,\n cross_replica=False, mybn=False,\n G_activation=nn.ReLU(inplace=False),\n G_lr=5e-5, G_B1=0.0, G_B2=0.999, adam_eps=1e-8,\n BN_eps=1e-5, SN_eps=1e-12, G_mixed_precision=False, G_fp16=False,\n G_init='ortho', skip_init=False, no_optim=False,\n G_param='SN', norm_style='bn',use_rangerlars=False,\n **kwargs):\n super(Generator, self).__init__()\n # Channel width mulitplier\n self.ch = G_ch\n # Dimensionality of the latent space\n self.dim_z = dim_z\n # The initial spatial dimensions\n self.bottom_width = bottom_width\n # Resolution of the output\n self.resolution = resolution\n # Kernel size?\n self.kernel_size = G_kernel_size\n # Attention?\n self.attention = G_attn\n # number of classes, for use in categorical conditional generation\n self.n_classes = n_classes\n # Use shared embeddings?\n self.G_shared = G_shared\n # Dimensionality of the shared embedding? Unused if not using G_shared\n self.shared_dim = shared_dim if shared_dim > 0 else dim_z\n # Hierarchical latent space?\n self.hier = hier\n # Cross replica batchnorm?\n self.cross_replica = cross_replica\n # Use my batchnorm?\n self.mybn = mybn\n # nonlinearity for residual blocks\n self.activation = G_activation\n # Initialization style\n self.init = G_init\n # Parameterization style\n self.G_param = G_param\n # Normalization style\n self.norm_style = norm_style\n # Epsilon for BatchNorm?\n self.BN_eps = BN_eps\n # Epsilon for Spectral Norm?\n self.SN_eps = SN_eps\n # fp16?\n self.fp16 = G_fp16\n # Architecture dict\n self.arch = G_arch(self.ch, self.attention)[resolution]\n\n # If using hierarchical latents, adjust z\n if self.hier:\n # Number of places z slots into\n self.num_slots = len(self.arch['in_channels']) + 1\n self.z_chunk_size = (self.dim_z // self.num_slots)\n # Recalculate latent dimensionality for even splitting into chunks\n self.dim_z = self.z_chunk_size * self.num_slots\n else:\n self.num_slots = 1\n self.z_chunk_size = 0\n\n # Which convs, batchnorms, and linear layers to use\n if self.G_param == 'SN':\n self.which_conv = partial(SNConv2d,\n kernel_size=3, padding=1,\n num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,\n eps=self.SN_eps)\n self.which_linear = partial(SNLinear,\n num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,\n eps=self.SN_eps)\n else:\n self.which_conv = partial(nn.Conv2d, kernel_size=3, padding=1)\n self.which_linear = nn.Linear\n\n # We use a non-spectral-normed embedding here regardless;\n # For some reason applying SN to G's embedding seems to randomly cripple G\n self.which_embedding = nn.Embedding\n bn_linear = (partial(self.which_linear, bias=False) if self.G_shared else self.which_embedding)\n self.which_bn = partial(ccbn,\n which_linear=bn_linear,\n cross_replica=self.cross_replica,\n mybn=self.mybn,\n input_size=(self.shared_dim + self.z_chunk_size if self.G_shared else self.n_classes),\n norm_style=self.norm_style,\n eps=self.BN_eps)\n\n # Prepare model\n # If not using shared embeddings, self.shared is just a passthrough\n self.shared = (self.which_embedding(n_classes, self.shared_dim) if G_shared else identity())\n # First linear layer\n self.linear = self.which_linear(\n self.dim_z // self.num_slots,\n self.arch['in_channels'][0] * (self.bottom_width ** 2))\n\n # self.blocks is a doubly-nested list of modules, the outer loop intended\n # to be over blocks at a given resolution (resblocks and/or self-attention)\n # while the inner loop is over a given block\n self.blocks = []\n for index in range(len(self.arch['out_channels'])):\n self.blocks += [[GBlock(\n in_channels=self.arch['in_channels'][index],\n out_channels=self.arch['out_channels'][index],\n which_conv=self.which_conv,\n which_bn=self.which_bn,\n activation=self.activation,\n upsample=(partial(F.interpolate, scale_factor=2) if self.arch['upsample'][index] else None)\n )]]\n\n # If attention on this block, attach it to the end\n if self.arch['attention'][self.arch['resolution'][index]]:\n print('Adding attention layer in G at resolution %d' % self.arch['resolution'][index])\n self.blocks[-1] += [Attention(self.arch['out_channels'][index], self.which_conv)]\n\n # Turn self.blocks into a ModuleList so that it's all properly registered.\n self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])\n\n # output layer: batchnorm-relu-conv.\n # Consider using a non-spectral conv here\n self.output_layer = nn.Sequential(\n bn(self.arch['out_channels'][-1], cross_replica=self.cross_replica, mybn=self.mybn),\n self.activation,\n self.which_conv(self.arch['out_channels'][-1], 3)\n )\n\n # Initialize weights. Optionally skip init for testing.\n if not skip_init:\n self.init_weights()\n\n # Set up optimizer\n # If this is an EMA copy, no need for an optim, so just return now\n if no_optim:\n return\n self.lr, self.B1, self.B2, self.adam_eps = G_lr, G_B1, G_B2, adam_eps\n if G_mixed_precision:\n print('Using fp16 adam in G...')\n self.optim = Adam16(params=self.parameters(), lr=self.lr,\n betas=(self.B1, self.B2), weight_decay=0,\n eps=self.adam_eps)\n\n elif use_rangerlars:\n self.optim=RangerLars(self.parameters(),lr=self.lr)\n else:\n self.optim = optim.Adam(params=self.parameters(), lr=self.lr,\n betas=(self.B1, self.B2), weight_decay=0,\n eps=self.adam_eps)\n\n # LR scheduling, left here for forward compatibility\n # self.lr_sched = {'itr' : 0}# if self.progressive else {}\n # self.j = 0\n\n # Initialize\n def init_weights(self):\n self.param_count = 0\n for module in self.modules():\n if (isinstance(module, nn.Conv2d)\n or isinstance(module, nn.Linear)\n or isinstance(module, nn.Embedding)):\n if self.init == 'ortho':\n init.orthogonal_(module.weight)\n elif self.init == 'N02':\n init.normal_(module.weight, 0, 0.02)\n elif self.init in ['glorot', 'xavier']:\n init.xavier_uniform_(module.weight)\n else:\n print('Init style not recognized...')\n self.param_count += sum([p.data.nelement() for p in module.parameters()])\n print('Param count for G''s initialized parameters: %d' % self.param_count)\n\n # Note on this forward function: we pass in a y vector which has\n # already been passed through G.shared to enable easy class-wise\n # interpolation later. If we passed in the one-hot and then ran it through\n # G.shared in this forward function, it would be harder to handle.\n def forward(self, z, y):\n # If hierarchical, concatenate zs and ys\n if self.hier:\n zs = torch.split(z, self.z_chunk_size, 1)\n z = zs[0]\n ys = [torch.cat([y, item], 1) for item in zs[1:]]\n else:\n ys = [y] * len(self.blocks)\n\n # First linear layer\n h = self.linear(z)\n # Reshape\n h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width)\n\n # Loop over blocks\n for index, blocklist in enumerate(self.blocks):\n # Second inner loop in case block has multiple layers\n for block in blocklist:\n h = block(h, ys[index])\n\n # Apply batchnorm-relu-conv-tanh at output\n return torch.tanh(self.output_layer(h))\n\n\n# Discriminator architecture, same paradigm as G's above\ndef D_arch(ch=64, attention='64', ksize='333333', dilation='111111'):\n arch = {}\n arch[256] = {'in_channels': [3] + [ch * item for item in [1, 2, 4, 8, 8, 16]],\n 'out_channels': [item * ch for item in [1, 2, 4, 8, 8, 16, 16]],\n 'downsample': [True] * 6 + [False],\n 'resolution': [128, 64, 32, 16, 8, 4, 4],\n 'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]\n for i in range(2, 8)}}\n arch[128] = {'in_channels': [3] + [ch * item for item in [1, 2, 4, 8, 16]],\n 'out_channels': [item * ch for item in [1, 2, 4, 8, 16, 16]],\n 'downsample': [True] * 5 + [False],\n 'resolution': [64, 32, 16, 8, 4, 4],\n 'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]\n for i in range(2, 8)}}\n arch[64] = {'in_channels': [3] + [ch * item for item in [1, 2, 4, 8]],\n 'out_channels': [item * ch for item in [1, 2, 4, 8, 16]],\n 'downsample': [True] * 4 + [False],\n 'resolution': [32, 16, 8, 4, 4],\n 'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]\n for i in range(2, 7)}}\n arch[32] = {'in_channels': [3] + [item * ch for item in [4, 4, 4]],\n 'out_channels': [item * ch for item in [4, 4, 4, 4]],\n 'downsample': [True, True, False, False],\n 'resolution': [16, 16, 16, 16],\n 'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]\n for i in range(2, 6)}}\n return arch\n\n\nclass Discriminator(nn.Module):\n\n def __init__(self, D_ch=64, D_wide=True, resolution=128,\n D_kernel_size=3, D_attn='64', n_classes=1000,\n num_D_SVs=1, num_D_SV_itrs=1, D_activation=nn.ReLU(inplace=False),\n D_lr=2e-4, D_B1=0.0, D_B2=0.999, adam_eps=1e-8,\n SN_eps=1e-12, output_dim=1, D_mixed_precision=False, D_fp16=False,\n D_init='ortho', skip_init=False, D_param='SN', **kwargs):\n super(Discriminator, self).__init__()\n # Width multiplier\n self.ch = D_ch\n # Use Wide D as in BigGAN and SA-GAN or skinny D as in SN-GAN?\n self.D_wide = D_wide\n # Resolution\n self.resolution = resolution\n # Kernel size\n self.kernel_size = D_kernel_size\n # Attention?\n self.attention = D_attn\n # Number of classes\n self.n_classes = n_classes\n # Activation\n self.activation = D_activation\n # Initialization style\n self.init = D_init\n # Parameterization style\n self.D_param = D_param\n # Epsilon for Spectral Norm?\n self.SN_eps = SN_eps\n # Fp16?\n self.fp16 = D_fp16\n # Architecture\n self.arch = D_arch(self.ch, self.attention)[resolution]\n\n # Which convs, batchnorms, and linear layers to use\n # No option to turn off SN in D right now\n if self.D_param == 'SN':\n self.which_conv = partial(SNConv2d,\n kernel_size=3, padding=1,\n num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,\n eps=self.SN_eps)\n self.which_linear = partial(SNLinear,\n num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,\n eps=self.SN_eps)\n self.which_embedding = partial(SNEmbedding,\n num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,\n eps=self.SN_eps)\n # Prepare model\n # self.blocks is a doubly-nested list of modules, the outer loop intended\n # to be over blocks at a given resolution (resblocks and/or self-attention)\n self.blocks = []\n for index in range(len(self.arch['out_channels'])):\n self.blocks += [[DBlock(\n in_channels=self.arch['in_channels'][index],\n out_channels=self.arch['out_channels'][index],\n which_conv=self.which_conv,\n wide=self.D_wide,\n activation=self.activation,\n preactivation=(index > 0),\n downsample=(nn.AvgPool2d(2) if self.arch['downsample'][index] else None)\n )]]\n # If attention on this block, attach it to the end\n if self.arch['attention'][self.arch['resolution'][index]]:\n print('Adding attention layer in D at resolution %d' % self.arch['resolution'][index])\n self.blocks[-1] += [Attention(self.arch['out_channels'][index], self.which_conv)]\n # Turn self.blocks into a ModuleList so that it's all properly registered.\n self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])\n # Linear output layer. The output dimension is typically 1, but may be\n # larger if we're e.g. turning this into a VAE with an inference output\n self.linear = self.which_linear(self.arch['out_channels'][-1], output_dim)\n # Embedding for projection discrimination\n self.embed = self.which_embedding(self.n_classes, self.arch['out_channels'][-1])\n\n # Initialize weights\n if not skip_init:\n self.init_weights()\n\n # Set up optimizer\n self.lr, self.B1, self.B2, self.adam_eps = D_lr, D_B1, D_B2, adam_eps\n if D_mixed_precision:\n print('Using fp16 adam in D...')\n self.optim = Adam16(params=self.parameters(), lr=self.lr,\n betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)\n else:\n self.optim = optim.Adam(params=self.parameters(), lr=self.lr,\n betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)\n # LR scheduling, left here for forward compatibility\n # self.lr_sched = {'itr' : 0}# if self.progressive else {}\n # self.j = 0\n\n # Initialize\n def init_weights(self):\n self.param_count = 0\n for module in self.modules():\n if (isinstance(module, nn.Conv2d)\n or isinstance(module, nn.Linear)\n or isinstance(module, nn.Embedding)):\n if self.init == 'ortho':\n init.orthogonal_(module.weight)\n elif self.init == 'N02':\n init.normal_(module.weight, 0, 0.02)\n elif self.init in ['glorot', 'xavier']:\n init.xavier_uniform_(module.weight)\n else:\n print('Init style not recognized...')\n self.param_count += sum([p.data.nelement() for p in module.parameters()])\n print('Param count for D''s initialized parameters: %d' % self.param_count)\n\n def forward(self, x, y=None):\n # Stick x into h for cleaner for loops without flow control\n h = x\n # Loop over blocks\n for index, blocklist in enumerate(self.blocks):\n for block in blocklist:\n h = block(h)\n # Apply global sum pooling as in SN-GAN\n h = torch.sum(self.activation(h), [2, 3])\n # Get initial class-unconditional output\n out = self.linear(h)\n # Get projection of final featureset onto class vectors and add to evidence\n out = out + torch.sum(self.embed(y) * h, 1, keepdim=True)\n return out\n\n\n# Parallelized G_D to minimize cross-gpu communication\n# Without this, Generator outputs would get all-gathered and then rebroadcast.\nclass G_D(nn.Module):\n def __init__(self, G, D):\n super(G_D, self).__init__()\n self.G = G\n self.D = D\n\n def forward(self, z, gy, x=None, dy=None, train_G=False, return_G_z=False,\n split_D=False):\n # If training G, enable grad tape\n with torch.set_grad_enabled(train_G):\n # Get Generator output given noise\n G_z = self.G(z, self.G.shared(gy))\n # Cast as necessary\n if self.G.fp16 and not self.D.fp16:\n G_z = G_z.float()\n if self.D.fp16 and not self.G.fp16:\n G_z = G_z.half()\n # Split_D means to run D once with real data and once with fake,\n # rather than concatenating along the batch dimension.\n if split_D:\n D_fake = self.D(G_z, gy)\n if x is not None:\n D_real = self.D(x, dy)\n return D_fake, D_real\n else:\n if return_G_z:\n return D_fake, G_z\n else:\n return D_fake\n # If real data is provided, concatenate it with the Generator's output\n # along the batch dimension for improved efficiency.\n else:\n D_input = torch.cat([G_z, x], 0) if x is not None else G_z\n D_class = torch.cat([gy, dy], 0) if dy is not None else gy\n # Get Discriminator output\n D_out = self.D(D_input, D_class)\n if x is not None:\n return torch.split(D_out, [G_z.shape[0], x.shape[0]]) # D_fake, D_real\n else:\n if return_G_z:\n return D_out, G_z\n else:\n return D_out\n\n\nclass Distribution(torch.Tensor):\n # Init the params of the distribution\n def init_distribution(self, dist_type, **kwargs):\n self.dist_type = dist_type\n self.dist_kwargs = kwargs\n if self.dist_type == 'normal':\n self.mean, self.var = kwargs['mean'], kwargs['var']\n elif self.dist_type == 'categorical':\n self.num_categories = kwargs['num_categories']\n\n def sample_(self):\n if self.dist_type == 'normal':\n self.normal_(self.mean, self.var)\n elif self.dist_type == 'categorical':\n self.random_(0, self.num_categories)\n # return self.variable\n\n # Silly hack: overwrite the to() method to wrap the new object\n # in a distribution as well\n def to(self, *args, **kwargs):\n new_obj = Distribution(self)\n new_obj.init_distribution(self.dist_type, **self.dist_kwargs)\n new_obj.data = super().to(*args, **kwargs)\n return new_obj\n\n\ndef prepare_z_y(G_batch_size, dim_z, nclasses, device='cuda', fp16=False, z_var=1.0):\n z_ = Distribution(torch.randn(G_batch_size, dim_z, requires_grad=False))\n z_.init_distribution('normal', mean=0, var=z_var)\n z_ = z_.to(device, torch.float16 if fp16 else torch.float32)\n\n if fp16:\n z_ = z_.half()\n\n y_ = Distribution(torch.zeros(G_batch_size, requires_grad=False))\n y_.init_distribution('categorical', num_categories=nclasses)\n y_ = y_.to(device, torch.int64)\n return z_, y_\n\n\n# Sample function for use with inception metrics\ndef sample_fn(G, z_, y_, config):\n with torch.no_grad():\n z_.sample_()\n y_.sample_()\n G_z = G(z_, G.shared(y_))\n return G_z, y_\n\n\ndef toggle_grad(model, on_or_off):\n for param in model.parameters():\n param.requires_grad = on_or_off\n\n\n# Simple wrapper that applies EMA to a model. COuld be better done in 1.0 using\n# the parameters() and buffers() module functions, but for now this works\n# with state_dicts using .copy_\nclass apply_ema(object):\n def __init__(self, source, target, decay=0.9999, start_itr=0):\n self.source = source\n self.target = target\n self.decay = decay\n # Optional parameter indicating what iteration to start the decay at\n self.start_itr = start_itr\n # Initialize target's params to be source's\n self.source_dict = self.source.state_dict()\n self.target_dict = self.target.state_dict()\n print('Initializing EMA parameters to be source parameters...')\n with torch.no_grad():\n for key in self.source_dict:\n self.target_dict[key].data.copy_(self.source_dict[key].data)\n # target_dict[key].data = source_dict[key].data # Doesn't work!\n\n def update(self, itr=None):\n # If an iteration counter is provided and itr is less than the start itr,\n # peg the ema weights to the underlying weights.\n if itr and itr < self.start_itr:\n decay = 0.0\n else:\n decay = self.decay\n with torch.no_grad():\n for key in self.source_dict:\n self.target_dict[key].data.copy_(self.target_dict[key].data * decay\n + self.source_dict[key].data*(1-decay))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device=torch.device('cuda:0')\n\nclass Mish(nn.Module):\n def __init__(self):\n super().__init__()\n\n def forward(self, x):\n #inlining this saves 1 second per epoch (V100 GPU) vs having a temp x and then returning x(!)\n return x *( torch.tanh(F.softplus(x)))\n\ndef loss_hinge_dis(dis_fake, dis_real):\n loss_real = torch.mean(F.relu(1. - dis_real))\n loss_fake = torch.mean(F.relu(1. + dis_fake))\n return loss_real, loss_fake\n\n\ndef loss_hinge_gen(dis_fake):\n loss = -torch.mean(dis_fake)\n return loss\n\ndef ortho(model, strength=1e-4, blacklist=[]):\n with torch.no_grad():\n for param in model.parameters():\n # Only apply this to parameters with at least 2 axes, and not in the blacklist\n if len(param.shape) < 2 or any([param is item for item in blacklist]):\n continue\n w = param.view(param.shape[0], -1)\n grad = (2 * torch.mm(torch.mm(w, w.t())\n * (1. - torch.eye(w.shape[0], device=w.device)), w))\n param.grad.data += strength * grad.view(param.shape)\n\ndef default_ortho(model, strength=1e-4, blacklist=[]):\n with torch.no_grad():\n for param in model.parameters():\n # Only apply this to parameters with at least 2 axes & not in blacklist\n if len(param.shape) < 2 or param in blacklist:\n continue\n w = param.view(param.shape[0], -1)\n grad = (2 * torch.mm(torch.mm(w, w.t())\n - torch.eye(w.shape[0], device=w.device), w))\n param.grad.data += strength * grad.view(param.shape)\n \ndef create_train_fn(G,D,GD,z_,y_,ema,state_dict,config):\n def train(x,y):\n G.optim.zero_grad()\n D.optim.zero_grad()\n\n x=torch.split(x,config['batch_size'])\n y=torch.split(y,config['batch_size'])\n counter=0\n\n if config['toggle_grads']:\n toggle_grad(D,True)\n toggle_grad(G,False)\n\n '''一个step内训练多次D'''\n \n for step_index in range(config['num_D_steps']):\n D.optim.zero_grad()\n \n for accumulation_index in range(config['num_D_accumulations']):\n z_.sample_()\n y_.sample_()\n \n D_fake, D_real = GD(z_[:config['batch_size']], y_[:config['batch_size']],\n x[counter], y[counter], train_G=False,\n split_D=config['split_D'])\n\n # Compute components of D's loss, average them, and divide by\n # the number of gradient accumulations\n D_loss_real, D_loss_fake = loss_hinge_dis(D_fake, D_real)\n D_loss = (D_loss_real + D_loss_fake) / float(config['num_D_accumulations'])\n D_loss.backward()\n counter += 1\n\n if config['D_ortho'] > 0.0:\n ortho(D, config['D_ortho'])\n\n D.optim.step()\n if config['toggle_grads']:\n toggle_grad(D, False)\n toggle_grad(G, True)\n\n # Zero G's gradients by default before training G, for safety\n '''一个step训练一次G'''\n G.optim.zero_grad()\n\n # If accumulating gradients, loop multiple times\n for accumulation_index in range(config['num_G_accumulations']):\n z_.sample_()\n y_.sample_()\n D_fake = GD(z_, y_, train_G=True, split_D=config['split_D'])\n G_loss = loss_hinge_gen(D_fake) / float(config['num_G_accumulations'])\n G_loss.backward()\n\n if config['G_ortho'] > 0.0:\n # Don't ortho reg shared, it makes no sense. Really we should blacklist any embeddings for this\n ortho(G, config['G_ortho'], blacklist=[param for param in G.shared.parameters()])\n G.optim.step()\n\n # If we have an ema, update it, regardless of if we test with it or not\n if config['ema']:\n ema.update(state_dict['itr'])\n\n out = {'G_loss': float(G_loss.item()),\n 'D_loss_real': float(D_loss_real.item()),\n 'D_loss_fake': float(D_loss_fake.item())}\n # Return G's loss and the components of D's loss.\n return out\n\n return train\n\ndef run(G, D, GD, loaders, state_dict, config, device, G_ema=None, ema=None):\n G_batch_size = max(config['G_batch_size'], config['batch_size'])\n '''随机抽取用于训练的噪声和label'''\n\n z_, y_ = prepare_z_y(\n G_batch_size, G.dim_z, config['n_classes'], device=device, fp16=config['G_fp16'])\n '''用于测试用的噪声和label'''\n fixed_z, fixed_y = prepare_z_y(\n G_batch_size, G.dim_z, config['n_classes'], device=device, fp16=config['G_fp16'])\n\n fixed_z.sample_()\n fixed_y.sample_()\n\n train=create_train_fn(G,D,GD,z_,y_,ema,state_dict,config)\n total_iters=config['num_epochs']*len(loaders)\n\n\n\n for epoch in range(state_dict['epoch'],config['num_epochs']):\n total_G_loss=0\n total_D_fake_loss=0\n total_D_real_loss=0\n for i,(x,y) in enumerate(loaders):\n state_dict['itr']+=1\n\n G.train()\n D.train()\n if config['ema']:\n G_ema.train()\n if config['D_fp16']:\n x,y=x.to(device).half(),y.to(device)\n else:\n x,y=x.to(device),y.to(device)\n\n '''得到这个step的损失'''\n metrics=train(x,y)\n total_G_loss+=metrics['G_loss']\n total_D_fake_loss+=metrics['D_loss_fake']\n total_D_real_loss+=metrics['D_loss_real']\n print('epoch:{},G_loss:{},D_fake_loss:{},D_real_loss:{}'.format(epoch,total_G_loss/len(loaders),total_D_fake_loss/len(loaders),total_D_real_loss/len(loaders)))\n state_dict['epoch'] += 1\n\n\nactivation_dict = {\n 'inplace_relu': nn.ReLU(inplace=True),\n 'relu': nn.ReLU(inplace=False),\n 'ir': nn.ReLU(inplace=True),\n 'mish':Mish(),\n}\nconfig = {\n 'num_workers': 1,\n 'pin_memory': False,\n 'shuffle': True,\n 'use_multiepoch_sampler': False,\n\n # Model\n 'model': 'BigGAN',\n 'G_param': 'SN',\n 'D_param': 'SN',\n 'G_ch': 64,\n 'D_ch': 64,\n 'G_depth': 1,\n 'D_depth': 1,\n 'D_wide': True,\n 'G_shared': False,\n 'shared_dim': 0,\n 'dim_z': 128,\n 'z_var': 1.0,\n 'hier': False,\n 'cross_replica': False,\n 'mybn': False,\n 'G_nl': 'relu',\n 'D_nl': 'relu',\n 'G_attn': '64',\n 'D_attn': '64',\n 'norm_style': 'bn',\n\n ### Model init stuff ###\n 'seed': 0,\n 'G_init': 'ortho',\n 'D_init': 'ortho',\n 'skip_init': False,\n\n ### Optimizer stuff ###\n 'G_lr': 5e-5,\n 'D_lr': 2e-4,\n 'G_B1': 0.0,\n 'D_B1': 0.0,\n 'G_B2': 0.999,\n 'D_B2': 0.999,\n\n ### Batch size, parallel, and precision stuff ###\n 'batch_size': 64,\n 'G_batch_size': 0,\n 'num_G_accumulations': 1,\n 'num_D_steps': 2,\n 'num_D_accumulations': 1,\n 'split_D': False,\n 'num_epochs': 100,\n 'parallel': False,\n 'G_fp16': False,\n 'D_fp16': False,\n 'D_mixed_precision': False,\n 'G_mixed_precision': False,\n 'accumulate_stats': False,\n 'num_standing_accumulations': 16,\n\n ### Bookkeping stuff ###\n 'G_eval_mode': False,\n 'save_every': 2000,\n 'num_save_copies': 2,\n 'num_best_copies': 2,\n 'which_best': 'IS',\n 'no_fid': False,\n 'test_every': 5000,\n 'num_inception_images': 50000,\n 'hashname': False,\n 'base_root': '',\n 'data_root': '../data',\n 'weights_root': '../weights',\n 'logs_root': '../logs',\n 'samples_root': '../samples',\n 'pbar': 'mine',\n 'name_suffix': '',\n 'experiment_name': '',\n 'config_from_name': False,\n\n ### EMA Stuff ###\n 'ema': False,\n 'ema_decay': 0.9999,\n 'use_ema': False,\n 'ema_start': 0,\n\n ### Numerical precision and SV stuff ###\n 'adam_eps': 1e-8,\n 'BN_eps': 1e-5,\n 'SN_eps': 1e-8,\n 'num_G_SVs': 1,\n 'num_D_SVs': 1,\n 'num_G_SV_itrs': 1,\n 'num_D_SV_itrs': 1,\n\n ### Ortho reg stuff ###\n 'G_ortho': 0.0, # 1e-4 is default for BigGAN\n 'D_ortho': 0.0,\n 'toggle_grads': True,\n\n ### Which train function ###\n 'which_train_fn': 'GAN',\n\n ### Resume training stuff\n 'load_weights': '',\n 'resume': False,\n\n ### Log stuff ###\n 'logstyle': '%3.3e',\n 'log_G_spectra': False,\n 'log_D_spectra': False,\n 'sv_log_interval': 10,\n 'log_interval': 100,\n}\nconfig['resolution'] = 64\nconfig['n_classes'] = 120\n\n# Current experiment\nconfig['num_epochs'] = 100\nconfig['num_D_steps'] = 1\nconfig['G_lr'] = 1e-4 * 0.85\nconfig['D_lr'] = 4e-4 * 0.85\nconfig['D_B2'] = 0.999\nconfig['G_B2'] = 0.999\nconfig['G_ch'] = 32\nconfig['D_ch'] = 64\nconfig['G_attn'] = '32'\nconfig['D_attn'] = '32'\nconfig['G_nl'] = 'inplace_relu'\nconfig['D_nl'] = 'inplace_relu'\nconfig['G_shared'] = True\nconfig['shared_dim'] = 128\nconfig['dim_z'] = 120\nconfig['hier'] = True\nconfig['SN_eps'] = 1e-8\nconfig['BN_eps'] = 1e-5\nconfig['adam_eps'] = 1e-8\nconfig['G_ortho'] = 0.0\nconfig['G_init'] = 'ortho'\nconfig['D_init'] = 'ortho'\nconfig['G_eval_mode'] = True\nconfig['save_every'] = 5000\nconfig['log_interval'] = 500\nconfig['batch_size'] = 64\n\nconfig['ema'] = True\nconfig['use_ema'] = True\nconfig['ema_decay'] = 0.9999\nconfig['ema_start'] = 50_000\n\nconfig['stop_after'] = 31_500\n\nconfig['trunc_z'] = 0.0\nconfig['denoise'] = True\nconfig['denoise_str_lum'] = 2.65\nconfig['denoise_str_col'] = 1.0\nconfig['denoise_kernel_size'] = 7\nconfig['denoise_search_window'] = 11\n\ndevice=torch.device('cuda:0')\nseed_rng(1234)\ntorch.backends.cudnn.benchmark=True\n\nD_batch_size=(config['batch_size'] * config['num_D_steps'] * config['num_D_accumulations'])\n\n\n# class_name = ['bear', 'bird', 'cat', 'dog', 'giraffe', 'horse', 'sheep', 'zebra']\n# class_dict = {}\n# data_dir = '/kaggle/input/myanimal/coco-animals'\n# for i, label in enumerate(class_name):\n# class_dict[label] = i\n\n# train_names = []\n# image_labels = {}\n# for i, label in enumerate(class_name):\n# for d in ['train', 'val']:\n# names = get_image_files(os.path.join(data_dir, d, label))\n# for name in names:\n# name = str(name)\n# train_names.append(name)\n# image_labels[name] = class_dict[label]\n\n# loaders=get_dataloader(train_names,image_labels,config['batch_size'])\nimport zipfile\nwith zipfile.ZipFile('/kaggle/input/generative-dog-images/all-dogs.zip',\"r\") as z:\n z.extractall(\".\")\n \npaths=[]\nfor root, _, fnames in sorted(os.walk('./all-dogs')):\n for fname in sorted(fnames):\n path = os.path.join(root, fname)\n paths.append(path)\nwith zipfile.ZipFile('/kaggle/input/generative-dog-images/Annotation.zip',\"r\") as z:\n z.extractall(\".\")\n \nINPUT_PATH = './all-dogs'\nANNOT_PATH = './Annotation'\nOUT_DIR = '../output_images'\n\nIMG_SIZE = 64\nIMG_SIZE_2 = IMG_SIZE * 2\nIMG_CHANNELS = 3\nIMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')\n\n\n# loaders=get_dogloader(paths,'./Annotation')\nloaders=get_data_loaders(INPUT_PATH,config['batch_size'])\nprint(len(loaders))\nconfig['G_activation'] = activation_dict[config['G_nl']]\nconfig['D_activation'] = activation_dict[config['D_nl']]\nG = Generator(**config).to(device)\nD = Discriminator(**config).to(device)\nif config['ema']:\n print('Preparing EMA for G with decay of {}'.format(config['ema_decay']))\n G_ema = Generator(**{**config, 'skip_init': True, 'no_optim': True}).to(device)\n ema = apply_ema(G, G_ema, config['ema_decay'], config['ema_start'])\nelse:\n G_ema, ema = None, None\n\nGD = G_D(G, D)\nstate_dict = {\n 'itr': 0,\n 'epoch': 0,\n 'config': config\n}\nrun(G, D, GD, loaders[0], state_dict, config, device, G_ema=G_ema, ema=ema)\n\n\n\n\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nwhich_G = G_ema if config['ema'] and config['use_ema'] else G\nif config['use_ema']:\n collect_bn_stats(G_ema, 10_000, config)\n \ngc.collect()\n \nG.eval()\nos.makedirs('users')\ngenerate_images('users',G,1000,config)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_saved_samples('users')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}