Github项目 - Human Image Gender Classifier
Github - Human Image Gender Classifier论文：Expressive Body Capture: 3D Hands, Face, and Body from a Single Image - CVPR2019主页：https://smpl-x.is.tue.mpg.de/基于OpenPose姿态估计结果，裁剪人体并识别人体性别.示例，如图：1. 项目环境Python 3.7TensorFlow 1.13.1ConfigerCUDA-10.0, cuDNN-7.5, Ubuntu 18.04requirements.txt:imageio==2.5.0 pillow numpy==1.16.3 scikit-image==0.15.0 scikit-learn==0.20.3 scipy==1.2.1 tensorflow==1.15.2安装：pip install git+https://github.com/nghorbani/homogenus模型下载，假设权重文件存放路径homogenus/trained_models/tf：https://smpl-x.is.tue.mpg.de/downloadsHomogenus: Human Image Gender Classifier2. 项目源码项目结构：├── homogenus │ ├── tf │ │ ├── homogenus_infer.py │ │ └── __init__.py │ ├── tools │ │ ├── body_cropper.py │ │ ├── image_tools.py │ │ ├── __init__.py │ │ └── omni_tools.py │ └── trained_models │ └── tf │ ├── checkpoint │ ├── TR02_E02_It_002010.ckpt.data-00000-of-00001 │ ├── TR02_E02_It_002010.ckpt.index │ ├── TR02_E02_It_002010.ckpt.meta │ └── version.txt ├── requirements.txt └── setup.py2.1. 模型预测 homogenus_infer.py#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- import tensorflow as tf import numpy as np import os, sys import json import cv2 import glob from homogenus.tools.image_tools import put_text_in_image, fontColors, read_prep_image, save_images from homogenus.tools.body_cropper import cropout_openpose, should_accept_pose from homogenus.tools.omni_tools import makepath class Homogenus_infer(object): def __init__(self, trained_model_dir, sess=None): ''' :param trained_model_dir: the directory where you have put the homogenus TF trained models :param sess: ''' best_model_fname = sorted(glob.glob(os.path.join(trained_model_dir , '*.ckpt.index')), key=os.path.getmtime) if len(best_model_fname): self.best_model_fname = best_model_fname[-1].replace('.index', '') else: raise ValueError('Couldnt find TF trained model in the provided directory --trained_model_dir=%s. Make sure you have downloaded them there.' % trained_model_dir) if sess is None: self.sess = tf.Session() else: self.sess = sess # Load graph. self.saver = tf.train.import_meta_graph(self.best_model_fname+'.meta') self.graph = tf.get_default_graph() self.prepare() def prepare(self): print('Restoring checkpoint %s..' % self.best_model_fname) self.saver.restore(self.sess, self.best_model_fname) def predict_genders(self, images_indir, openpose_indir, images_outdir=None, openpose_outdir=None): ''' Given a directory with images and another directory with corresponding openpose genereated jsons will augment openpose jsons with gender labels. :param images_indir: Input directory of images with common extensions :param openpose_indir: Input directory of openpose jsons :param images_outdir: If given will overlay the detected gender on detected humans that pass the criteria :param openpose_outdir: If given will dump the gendered openpose files in this directory. if not will augment the origianls :return: ''' sys.stdout.write('\nRunning homogenus on --images_indir=%s --openpose_indir=%s\n'%(images_indir, openpose_indir)) im_fnames = [] for img_ext in ['png', 'jpg', 'jpeg', 'bmp']: im_fnames.extend(glob.glob(os.path.join(images_indir, '*.%s'%img_ext))) if len(im_fnames): sys.stdout.write('Found %d images\n' % len(im_fnames)) else: raise ValueError('No images could be found in %s'%images_indir) accept_threshold = 0.9 crop_margin = 0.08 if images_outdir is not None: makepath(images_outdir) if openpose_outdir is None: openpose_outdir = openpose_indir else: makepath(openpose_outdir) Iph = self.graph.get_tensor_by_name(u'input_images:0') probs_op = self.graph.get_tensor_by_name(u'probs_op:0') for im_fname in im_fnames: im_basename = os.path.basename(im_fname) img_ext = im_basename.split('.')[-1] openpose_in_fname = os.path.join(openpose_indir, im_basename.replace('.%s'%img_ext, '_keypoints.json')) with open(openpose_in_fname, 'r') as f: pose_data = json.load(f) im_orig = cv2.imread(im_fname, 3)[:,:,::-1].copy() for opnpose_pIdx in range(len(pose_data['people'])): pose_data['people'][opnpose_pIdx]['gender_pd'] = 'neutral' pose = np.asarray(pose_data['people'][opnpose_pIdx]['pose_keypoints_2d']).reshape(-1, 3) if not should_accept_pose(pose, human_prob_thr=0.5): continue crop_info = cropout_openpose(im_fname, pose, want_image=True, crop_margin=crop_margin) cropped_image = crop_info['cropped_image'] if cropped_image.shape[0] < 200 or cropped_image.shape[1] < 200: continue img = read_prep_image(cropped_image)[np.newaxis] probs_ob = self.sess.run(probs_op, feed_dict={Iph: img})[0] gender_id = np.argmax(probs_ob, axis=0) gender_prob = probs_ob[gender_id] gender_pd = 'male' if gender_id == 0 else 'female' if gender_prob>accept_threshold: color = 'green' text = 'pred:%s[%.3f]' % (gender_pd, gender_prob) else: text = 'thr:%s_pred:%s[%.3f]' % ('neutral', gender_pd, gender_prob) gender_pd = 'neutral' color = 'grey' x1 = crop_info['crop_boundary']['offset_width'] y1 = crop_info['crop_boundary']['offset_height'] x2 = crop_info['crop_boundary']['target_width'] + x1 y2 = crop_info['crop_boundary']['target_height'] + y1 im_orig = cv2.rectangle(im_orig, (x1, y1), (x2, y2), fontColors[color], 2) im_orig = put_text_in_image(im_orig, [text], color, (x1, y1))[0] pose_data['people'][opnpose_pIdx]['gender_pd'] = gender_pd sys.stdout.write('%s -- peron_id %d --> %s\n'%(im_fname, opnpose_pIdx, gender_pd)) if images_outdir != None: save_images(im_orig, images_outdir, [os.path.basename(im_fname)]) openpose_out_fname = os.path.join(openpose_outdir, im_basename.replace('.%s'%img_ext, '_keypoints.json')) with open(openpose_out_fname, 'w') as f: json.dump(pose_data, f) if images_outdir is not None: sys.stdout.write('Dumped overlayed images at %s'%images_outdir) sys.stdout.write('Dumped gendered openpose keypoints at %s'%openpose_outdir) if __name__ == '__main__': import argparse parser = argparse.ArgumentParser() parser.add_argument("-tm", "--trained_model_dir", default="./homogenus/trained_models/tf/", help="The path to the directory holding homogenus trained models in TF.") parser.add_argument("-ii", "--images_indir", required= True, help="Directory of the input images.") parser.add_argument("-oi", "--openpose_indir", required=True, help="Directory of openpose keypoints, e.g. json files.") parser.add_argument("-io", "--images_outdir", default=None, help="Directory to put predicted gender overlays. If not given, wont produce any overlays.") parser.add_argument("-oo", "--openpose_outdir", default=None, help="Directory to put the openpose gendered keypoints. If not given, it will augment the original openpose json files.") ps = parser.parse_args() hg = Homogenus_infer(trained_model_dir=ps.trained_model_dir) hg.predict_genders(images_indir=ps.images_indir, openpose_indir=ps.openpose_indir, images_outdir=ps.images_outdir, openpose_outdir=ps.openpose_outdir)2.2. body_cropper.py#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- from PIL import Image import numpy as np import os import json from homogenus.tools.image_tools import cropout_openpose def should_accept_pose(pose, human_prob_thr=.5): ''' :param pose: :param human_prob_thr: :return: ''' rleg_ids = [12,13] lleg_ids = [9,10] rarm_ids = [5,6,7] larm_ids = [2,3,4] head_ids = [16,15,14,17,0,1] human_prob = pose[:, 2].mean() rleg = sum(pose[rleg_ids][:, 2] > 0.0) lleg = sum(pose[lleg_ids][:, 2] > 0.0) rarm = sum(pose[rarm_ids][:, 2] > 0.0) larm = sum(pose[larm_ids][:, 2] > 0.0) head = sum(pose[head_ids][:, 2] > 0.0) if rleg<1 and lleg<1: return False if rarm<1 and larm<1: return False if head<2: return False if human_prob < human_prob_thr: return False return True def crop_humans(im_fname, pose_fname, want_image=True, human_prob_thr=0.5): ''' :param im_fname: the input image path :param pose_fname: the corresponding openpose json file :param want_image: if False will only return the crop boundaries, otherwise will also return the cropped image :param human_prob_thr: the probability to threshold the detected humans :return: ''' crop_infos = {} with open(pose_fname) as f: pose_data = json.load(f) if not len(pose_data['people']): return crop_infos for pIdx in range(len(pose_data['people'])): pose = np.asarray(pose_data['people'][pIdx]['pose_keypoints_2d']).reshape(-1, 3) if not should_accept_pose(pose, human_prob_thr=human_prob_thr): continue crop_info = cropout_openpose(im_fname, pose, want_image=want_image) crop_infos['%02d'%pIdx] = { 'crop_info':crop_info['crop_boundary'], #'pose':pose, 'pose_hash':hash(pose.tostring())} if want_image: crop_infos['%02d' % pIdx]['cropped_image'] = crop_info['cropped_image'].astype(np.uint8) return crop_infos def crop_dataset(base_dir, dataset_name, human_prob_thr=.5): import glob import random results_dir = os.path.join(base_dir, dataset_name, 'cropped_body_tight') images_dir = os.path.join(base_dir, dataset_name, 'images') pose_jsonpath = os.path.join(base_dir, dataset_name, 'openpose_json') if not os.path.exists(results_dir): os.makedirs(results_dir) crop_infos_jsonpath = os.path.join(results_dir, 'crop_infos.json') fnames = glob.glob(os.path.join(images_dir, '*.jpg')) random.shuffle(fnames) crop_infos = {} for fname in fnames: pose_fname = os.path.join(pose_jsonpath, os.path.basename(fname).replace('.jpg', '_keypoints.json')) if not os.path.exists(pose_fname): continue cur_crop_info = crop_humans(fname, pose_fname, human_prob_thr=human_prob_thr, want_image=True) for pname in cur_crop_info.keys(): crop_id = '%s_%s' % (os.path.basename(fname).split('.')[0], pname) crop_outpath = os.path.join(results_dir, '%s.jpg' % crop_id) cropped_image = cur_crop_info[pname]['cropped_image'] if cropped_image.shape[0] < 200 or cropped_image.shape[1] < 200: continue cur_crop_info.pop('cropped_image') # print(crop_outpath) result = Image.fromarray(cropped_image[:,:,::-1]) result.save(crop_outpath) crop_infos[crop_id] = cur_crop_info with open(crop_infos_jsonpath, 'w') as f: json.dump(crop_infos, f)2.3. image_tools.py#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- import cv2 import numpy as np import os fontColors = {'red': (255, 0, 0), 'green': (0, 255, 0), 'yellow': (255, 255, 0), 'blue': (0, 255, 255), 'orange': (255, 165, 0), 'black': (0, 0, 0), 'grey': (169, 169, 169), 'white': (255, 255, 255), } def crop_to_bounding_box(image, offset_height, offset_width, target_height, target_width): cropped = image[offset_height:offset_height + target_height, offset_width:offset_width + target_width, :] return cropped def pad_to_bounding_box(image, offset_height, offset_width, target_height, target_width): height, width, depth = image.shape after_padding_width = target_width - offset_width - width after_padding_height = target_height - offset_height - height # Do not pad on the depth dimensions. paddings = ((offset_height, after_padding_height), (offset_width, after_padding_width), (0, 0)) padded = np.pad(image, paddings, 'constant') return padded def resize_image_with_crop_or_pad(image, target_height, target_width): # crop to ratio, center height, width, c = image.shape width_diff = target_width - width offset_crop_width = max(-width_diff // 2, 0) offset_pad_width = max(width_diff // 2, 0) height_diff = target_height - height offset_crop_height = max(-height_diff // 2, 0) offset_pad_height = max(height_diff // 2, 0) # Maybe crop if needed. # print('image shape', image.shape) cropped = crop_to_bounding_box(image, offset_crop_height, offset_crop_width, min(target_height, height), min(target_width, width)) # print('after cropp', cropped.shape) # Maybe pad if needed. resized = pad_to_bounding_box(cropped, offset_pad_height, offset_pad_width, target_height, target_width) # print('after pad', resized.shape) return resized[:target_height, :target_width, :] def cropout_openpose(im_fname, pose, want_image=True, crop_margin=0.08): im_orig = cv2.imread(im_fname, 3) im_height, im_width = im_orig.shape[0], im_orig.shape[1] pose = pose[pose[:, 2] > 0.0] x_min, x_max = pose[:, 0].min(), pose[:, 0].max() y_min, y_max = pose[:, 1].min(), pose[:, 1].max() margin_h = crop_margin * im_height margin_w = crop_margin * im_width offset_height = int(max((y_min - margin_h), 0)) target_height = int(min((y_max + margin_h), im_height)) - offset_height offset_width = int(max((x_min - margin_w), 0)) target_width = int(min((x_max + margin_w), im_width)) - offset_width crop_info = {'crop_boundary': {'offset_height':offset_height, 'target_height':target_height, 'offset_width':offset_width, 'target_width':target_width}} if want_image: crop_info['cropped_image'] = crop_to_bounding_box(im_orig, offset_height, offset_width, target_height, target_width) return crop_info def put_text_in_image(images, text, color ='white', position=None): ''' :param images: 4D array of images :param text: list of text to be printed in each image :param color: the color or colors of each text :return: ''' import cv2 if not isinstance(text, list): text = [text] if not isinstance(color, list): color = [color for _ in range(images.shape[0])] if images.ndim == 3: images = images.reshape(1,images.shape[0],images.shape[1],3) images_out = [] for imIdx in range(images.shape[0]): img = images[imIdx].astype(np.uint8) font = cv2.FONT_HERSHEY_SIMPLEX if position is None:position = (10, img.shape[1]) fontScale = 1. lineType = 2 fontColor = fontColors[color[imIdx]] cv2.putText(img, text[imIdx], position, font, fontScale, fontColor, lineType) images_out.append(img) return np.array(images_out) def read_prep_image(im_fname, avoid_distortion=True): ''' if min(height, width) is larger than 224 subsample to 224. this will also affect the larger dimension. in the end crop and pad the whole image to get to 224x224 :param im_fname: :return: ''' if isinstance(im_fname, np.ndarray): image_data = im_fname else: image_data = cv2.imread(im_fname, 3) # height, width = image_reader.read_image_dims(sess, image_data) # image_data = image_reader.decode_jpeg(sess, image_data) # print(image_data.min(), image_data.max(), image_data.shape) # import matplotlib.pyplot as plt # plt.imshow(image_data[:,:,::-1].astype(np.uint8)) # plt.show() # height, width = image_data.shape[0], image_data.shape[1] # if min(height, width) > 224: # print(image_data.shape) # rt = 224. / min(height, width) # image_data = cv2.resize(image_data, (int(rt * width), int(rt * height)), interpolation=cv2.INTER_AREA) # print('>>resized to>>',image_data.shape) height, width = image_data.shape[0], image_data.shape[1] if avoid_distortion: if max(height, width) > 224: # print(image_data.shape) rt = 224. / max(height, width) image_data = cv2.resize(image_data, (int(rt * width), int(rt * height)), interpolation=cv2.INTER_AREA) # print('>>resized to>>',image_data.shape) else: from skimage.transform import resize image_data = resize(image_data, (224, 224), mode='constant', anti_aliasing=False, preserve_range=True) # print(image_data.min(), image_data.max(), image_data.shape) # import matplotlib.pyplot as plt # plt.imshow(image_data[:,:,::-1].astype(np.uint8)) # plt.show() image_data = resize_image_with_crop_or_pad(image_data, 224, 224) # print(image_data.min(), image_data.max(), image_data.shape) # import matplotlib.pyplot as plt # plt.imshow(image_data[:, :, ::-1].astype(np.uint8)) # plt.show() #return image_data.astype(np.float32) return image_data.astype(np.uint8) def save_images(images, out_dir, im_names = None): from homogenus.tools.omni_tools import id_generator if images.ndim == 3: images = images.reshape(1,images.shape[0],images.shape[1],3) from PIL import Image if im_names is None: im_names = ['%s.jpg'%id_generator(4) for i in range(images.shape[0])] for imIdx in range(images.shape[0]): result = Image.fromarray(images[imIdx].astype(np.uint8)) result.save(os.path.join(out_dir, im_names[imIdx])) return True2.4. omni_tools.py#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- import numpy as np copy2cpu = lambda tensor: tensor.detach().cpu().numpy() colors = { 'pink': [.7, .7, .9], 'purple': [.9, .7, .7], 'cyan': [.7, .75, .5], 'red': [1.0,0.0,0.0], 'green': [.0, 1., .0], 'yellow': [1., 1., 0], 'brown': [.5, .7, .7], 'blue': [.0, .0, 1.], 'offwhite': [.8, .9, .9], 'white': [1., 1., 1.], 'orange': [.5, .65, .9], 'grey': [.7, .7, .7], 'black': np.zeros(3), 'white': np.ones(3), 'yellowg': [0.83,1,0], } def id_generator(size=13): import string import random chars = string.ascii_uppercase + string.digits return ''.join(random.choice(chars) for _ in range(size)) def log2file(logpath=None, auto_newline = True): import sys if logpath is not None: makepath(logpath, isfile=True) fhandle = open(logpath,'a+') else: fhandle = None def _(text): if auto_newline: if not text.endswith('\n'): text = text + '\n' sys.stderr.write(text) if fhandle is not None: fhandle.write(text) fhandle.flush() return lambda text: _(text) def makepath(desired_path, isfile = False): ''' if the path does not exist make it :param desired_path: can be path to a file or a folder name :return: ''' import os if isfile: if not os.path.exists(os.path.dirname(desired_path)): os.makedirs(os.path.dirname(desired_path)) else: if not os.path.exists(desired_path): os.makedirs(desired_path) return desired_path
TensorFlow Serving 是一个适用于机器学习模型的灵活、高性能应用系统，专为生产环境而设计.借助 TensorFlow Serving，可以轻松部署新算法和实验，同时保留相同的服务器架构和 API.TensorFlow Serving 提供与 TensorFlow 模型的开箱即用型集成，但也可以轻松扩展以应用其他类型的模型和数据.TensorFlow Serving 可以将训练好的机器学习模型部署到线上，使用 gRPC 作为接口接受外部调用.TensorFlow Serving 支持模型热更新与自动模型版本管理。这意味着一旦部署 TensorFlow Serving 后，再也不需要为线上服务操心，只需要关心线下模型训练.https://www.tensorflow.org/tfx/serving/architecture1. 核心概念1.1. ServablesServables 是 TF Serving 的中心抽象.Servables 是客户端用来执行计算的底层对象，如查找和推断.Servables 的尺寸和粒度是灵活的. 单个 Servable 可以是任何东西，从查询表的一个分片，到推断模型元组的单个模型.Servables 可以是任意类型和任意推断，以确保灵活性和将来的进一步提升，如：流结果试验性 APIs异步操作Servables 并不管理其自己的生命周期.Servabels 一般包含如下：TensorFlow SavedModelBundle (tensorflow::Session)用于嵌入或词汇查询的查询表1.2. Servable VersionsTF Serving 可以在耽搁服务器实例的生命周期内，处理一个或多个版本的 servable. 其有助于随着时间更新加载新的算法配置、权重和其他数据. Servable Versions 支持同时加载多个版本的 servable，支持逐步展开和试验.在使用服务时，客户端可以请求特定的 version id 或最新版本的特定模型.1.3. Servable StreamsServable Streams 是 servable 版本的徐蕾，根据版本号增量排序.1.4. ModelsTF Serving 将 model 表示为一个或多个 servable. 一个机器学习模型包含一个或多个算法(包括学习的权重) 和查询表(或嵌入表).组合模型(composite model) 可以是如下情形之一：多个独立的 servables单个组合的 servable此外，一个 servable 还可以对应于一个模型的一部分. 例如，比较大的查询表可以跨很多 TF Serving 实例.1.5. LoadersLoaders 管理 servable 的生命周期.Loader API 支持独立于特定学习算法、数据或产品用例的公共基础设施.Loader API 标准化了用于 servable 的加载和卸载的 API.1.6. SourcesSources 是用于寻找和提供 servable 的插件模块.每个 Source 提供 0 个或多个 servable streams. 对于每个 servable stream，一个 Source 为每个可供加载的 version，提供一个 Loader 实例.TF Serving 的 Sources 接口可以用于从任意存储系统中发现 servables. TF Serving 内置了通用参考 Source 实现. 例如，Sources 能够访问诸如 RPC 之类的机制，且可以轮询文件系统.Sources 能够维护跨多个 servables 或版本所共享的状态. 其有助于使用 delta(diff) 更新不同 versions 间的 servables.1.7. Aspired VersionsAspired Versions 表示应该被加载和准备好了的 servable versions 的集合.Sources 每次为单个 servable stream 传递该 servable visions 集合.当一个 Source 向 Manager 给定新的 aspired versions 列表时，其将取代 servable stream 的现钱版本. Manager 将卸载列表中不再显示的所有以前加载的版本.1.8. ManagersManagers 处理 Servables 的完整生命周期，包括：加载 Servables部署 Servables 服务卸载 ServablesManagers 监听 Sources 并追踪所有的 versions.Managers 尝试完成 Sources 的请求，但如果所需的资源不可用时，则会拒绝加载 aspired version.Managers 也可以推迟卸载，例如，Manage 可能会等待卸载，直到新版本的完成加载. 这是基于确保始终至少加载一个版本的策略.TF Serving Manager 提供了简单直接的接口 - GetServableHandle()，用于客户端访问加载的 servable 实例.1.9. Core采用标准的 TF Serving APIs, TF Serving Core 管理 servables 的几个方面：生命周期 lifecycle度量 metrics2. Servable 生命如图，广义的来说，[1] - Sources 创建 Servable Versions 的 Loaders.[2] - Loaders 被作为 Aspired Versions 送到 Manager，以加载并部署服务到客户端请求.更细节的来说，[1] - Source 插件创建特定 version 的 Loader. 该 Loader 包含了 Servable 加载所需的元数据.[2] - Source 采用回调(callback) 方式通知 Aspired Version 的 Manager.[3] - Manager 采用配置的 Version Policy 来判断下一步要采取的行动，其可以卸载先前的加载版本或新增版本.[4] - 如果 Manager 判断其是安全的，则给到 Loader 所需的资源，并告诉 Loader 加载新版本.[5] - 客户端Clients 查询 Manager 以得到 Servable，或者是显式的指定版本，或者是请求最新版本. Manager 返回 Servable 的句柄(handle).例如，假设一个 Source 表示一个频繁更新模型权重的 TF graph. 其权重存储在磁盘的文件里.[1] - Source 检测到模型权重的一个新版本，其会创建一个包含指向磁盘模型数据的指针的 Loader.[2] - Source 通知 Aspired Version 的 Dynamic Manager.[3] - Dynamic Manager 根据 Version Policy，加载新版本模型.[4] - Dynamic Manager 通知 Loader 内存足够. Loader 采用新的权重实例化 TF graph.[5] - 客户端向模型的最新版本请求一个句柄，Dynamic Manager 返回 Servable 新版本的句柄.大致思路如下：[1] - 首先export model，准备计算图的代码，将保存的参数restore进入计算图；[2] - 然后调用export model的方法，导出模型[3] - 模型导出之后就直接拉取tensorflow serving的镜像然后创建，这里注意如果是tensorflow-serving-1.x.0的镜像，docker run之后就直接可以访问容器服务，tensorflow-serving-1.x.0-devel的镜像需要docker run之后，再进入容器运行.[4] - tensorflow_model_server指定容器映射的端口，model_name等等，如果计算图没有其他依赖，直接拉取tensorflow-serving-1.x.0就行[5] - 创建tensorflow serving镜像时需要指定端口映射，然后可以通过这个端口访问serving服务，还需要将宿主机自身模型的path挂载到容器的/models上[6] - 如何使用运行好的容器服务？可以通过RESTful API访问，也可以通过gRPC API访问[7] - 最后模型部分已经全部交给了tensorflow-server，再去写一个web服务包装一下就可以快速部署了.
原文：Latent Space Exploration with StyleGAN2 - 2020.08.10主要介绍 StyleGAN2 中潜空间（latent space）是什么，以及其能做什么.StyleGAN2关于 latent space + StyleGAN2 的博文：Latent Space Understanding - Ekin Tiu - Understanding Latent Space in Machine LearningStyleGAN2 结构 - Connor Shorten - StyleGAN2GANs 概述 + StyleGAN2 改变了什么 - akira - From GAN basic to StyleGAN2注：https://github.com/AmarSaini/Epoching_StyleGan2_Setup本文基于 TensorFlow 1.14本文不过多关注于模型，采用官方的 StyleGan2 Repo本文进行的结构：[0] Set-Up 只运行一次必要文件拉取[1] - GANs 概览回顾 z latent vector、Generator 和 DiscriminatorGANs vs VAEs[2] - 生成人脸采用官方 StyleGAN2 项目创建 Generator 输出可视化生成输出的 latent codes.[3] - Latent Codes 插值采用 Generator 输出的 latent codes 得到人脸图像[4] - 采用 landmark detection 方法进行人脸图像对齐对齐（归一化）输入图像，以进行 latent space 投影（projection）[5] - 将输入图像投影到 Latent Space学习对齐后输入图像的 latent codes投影的 latent codes 的插值[6] - Latent Directions/Controls 修改投影图像采用预先计算的 latent directions 来改变输入图像的人脸特征[7] - Widget-App 交互0. Set-Up下载项目：git clone https://github.com/AmarSaini/Epoching_StyleGan2_Setup.git项目处理：import shutil from pathlib import Path repo_root = Path('Epoching_StyleGan2_Setup/') # Pull contents out of the repo, into our current directory. for content in repo_root.iterdir(): shutil.move(str(content), '.') shutil.rmtree(repo_root)安装python依赖库：pip install requests pip install Pillow pip install tqdm pip install dlib pip install tensorflow-gpu==1.141. GANs 概述GANs，Generative Adversarial Networks，包含两个模型：[1] - Generator模型，将 latent code 转换为某种输出，如，人脸图像.[2] - Discriminator模型，判断某种输入，如，人脸图像，是真是假.Real: 图像来自于数据集Fake: 图像来自于 GeneratorGenerator 的输入是 latent code z，其可以是任意期望长度的向量，如，512维的向量.训练时，latent code 是随机采样的，512 维的随机向量当 latent code 随机采样时，称其为 latent random variable神奇的 latent code 存储的信息，使得 Generator 生成特定的输出如果可以找到特定输入的 latent code，则可以将其采用更少的数据来表示（如，将某人的图像表示为仅为 512 长度的 latent vector，相比于原始图像大小会小很多）.注：不要混淆了 GANs 和 VAEs, Variational Auto-EncodersGANs 学习从 random latent vectors 中生成输出，其模型输入数据的分布；但不必须是输入数据的精确样本VAEs 将输入样本编码为 latent vectors，并学习将 latent vectors 解码为与原始样本最接近的形式GANs 与 VAEs 的最大不同是，GANs 的 Generator 从未看到过输入图像，因此，无法将图像自动转换为对应的 latent code. （即，投影 projection ）From: https://www.slideshare.net/xavigiro/deep-learning-for-computer-vision-generative-models-and-adversarial-training-upc-20162. 生成人脸import sys sys.path.append('stylegan2/') from stylegan2 import pretrained_networks from stylegan2 import dnnlib from stylegan2.dnnlib import tflib from pathlib import Path from PIL import Image import pickle import numpy as np from tqdm import tqdm model_path = 'gdrive:networks/stylegan2-ffhq-config-f.pkl' fps = 20 results_size = 400加载 StyleGAN2 模型：# Code to load the StyleGAN2 Model def load_model(): _G, _D, Gs = pretrained_networks.load_networks(model_path) noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')] Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) Gs_kwargs.randomize_noise = False return Gs, noise_vars, Gs_kwargs # Generate images given a random seed (Integer) def generate_image_random(rand_seed): rnd = np.random.RandomState(rand_seed) z = rnd.randn(1, *Gs.input_shape[1:]) tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) images = Gs.run(z, None, **Gs_kwargs) return images, z # Generate images given a latent code ( vector of size [1, 512] ) def generate_image_from_z(z): images = Gs.run(z, None, **Gs_kwargs) return images然后，随机生成一些输出：Gs, noise_vars, Gs_kwargs = load_model() images, latent_code1 = generate_image_random(42) image1 = Image.fromarray(images[0]).resize((results_size, results_size)) print(latent_code1.shape) #(1, 512) image1.show()images, latent_code2 = generate_image_random(1234) image2 = Image.fromarray(images[0]).resize((results_size, results_size)) print(latent_code2.shape) #(1, 512) print(latent_code1[0][:5], latent_code2[0][:5]) #(array([ 0.49671415, -0.1382643 , 0.64768854, 1.52302986, -0.23415337]), # array([ 0.47143516, -1.19097569, 1.43270697, -0.3126519 , -0.72058873])) image2.show()3. Latent Codes 插值latent space 是某些数据的压缩表示，外形相似的样本在 latent space 应该是互相接近的.如果 latent space 比较好，则可以在 latent space 中的点进行过渡/插值（transition/interpolate），以得到中间的输出值.换句话说，可以对两张人脸进行变形. 如图：From：https://hackernoon.com/latent-space-visualization-deep-learning-bits-2-bd09a46920df如下对 latent_code1 和 latent_code2 进行插值：def linear_interpolate(code1, code2, alpha): return code1 * alpha + code2 * (1 - alpha) interpolated_latent_code = linear_interpolate(latent_code1, latent_code2, 0.5) print(interpolated_latent_code.shape) #(1, 512) images = generate_image_from_z(interpolated_latent_code) image_inter = Image.fromarray(images[0]).resize((results_size, results_size)) image_inter.show()动态化该过程：output_gifs_path = Path('output_gifs') # Make Output Gifs folder if it doesn't exist. if not output_gifs_path.exists(): output_gifs_path.mkdir() def get_concat_h(im1, im2): dst = Image.new('RGB', (im1.width + im2.width, im1.height)) dst.paste(im1, (0, 0)) dst.paste(im2, (im1.width, 0)) return dst def make_latent_interp_animation(code1, code2, img1, img2, num_interps): step_size = 1.0/num_interps all_imgs = [] amounts = np.arange(0, 1, step_size) for alpha in tqdm(amounts): interpolated_latent_code = linear_interpolate(code1, code2, alpha) images = generate_image_from_z(interpolated_latent_code) interp_latent_image = Image.fromarray(images[0]).resize((400, 400)) frame = get_concat_h(img1, interp_latent_image) frame = get_concat_h(frame, img2) all_imgs.append(frame) # save_name = output_gifs_path/'latent_space_traversal.gif' all_imgs[0].save(save_name, save_all=True, append_images=all_imgs[1:], duration=1000/fps, loop=0) # make_latent_interp_animation(latent_code1, latent_code2, image1, image2, num_interps=200) 如：4. 人脸图像对齐对于自己的人脸图像，而不是随机生成的人脸图像，如何进行投影到 latent space？首先，采用人脸关键点检测模型，对齐（归一化）图像.如，orig_img_path = Path('imgs') aligned_imgs_path = Path('aligned_imgs') # Make Aligned Images folder if it doesn't exist. if not aligned_imgs_path.exists(): aligned_imgs_path.mkdir() print(orig_img_path, aligned_imgs_path) #if not Path('shape_predictor_68_face_landmarks.dat').exists(): # wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2 # bzip2 -dv shape_predictor_68_face_landmarks.dat.bz2 from align_face import align_face # Align all of our images using a landmark detection model! all_imgs = list(orig_img_path.iterdir()) for img in all_imgs: align_face(str(img)).save(aligned_imgs_path/('aligned_'+img.name)) # aligned_img_set = list(aligned_imgs_path.iterdir()) aligned_img_set.sort() aligned_img_set = [Image.open(x) for x in aligned_img_set] orig_img_set = list(orig_img_path.iterdir()) orig_img_set.sort() orig_img_set = [Image.open(x) for x in orig_img_set] #对齐后的图像，转化为输出大小为 1024x1024 get_concat_h(orig_img_set[0], aligned_img_set[0]) get_concat_h(orig_img_set[1], aligned_img_set[1])如：5. 投影图像到 Latent Space#GPU上运行 import sys sys.path.append('stylegan2/') from stylegan2 import pretrained_networks from stylegan2 import dnnlib from stylegan2.dnnlib import tflib from pathlib import Path from PIL import Image import pickle import numpy as np from tqdm import tqdm model_path = 'gdrive:networks/stylegan2-ffhq-config-f.pkl' fps = 20 results_size = 400将数据转换为 StyleGAN2 兼容格式，python -W ignore stylegan2/dataset_tool.py create_from_images datasets_stylegan2/custom_imgs aligned_imgs/注：[1] - 将图像投影到 latent space，意味着，可视化出 latent code 使得 Generator 生成的输出如何与输入图像相似[2] - 问题是，如何可视化出 latent code？ VAEs 模型中，可以将图像送入解码器，即可得到 latent code.[3] - GANs 模型中，没有直接的方式来提取输入图像的 latent codes，但是可以进行优化得到.对于迭代过程中，主要进行：[1] - 用 Generator 从起始 latent vector 生成一些输出[2] - 得到 Generator 的输出图像，以及目标图像，共同送入 VGG16 模型（图像特征提取器）[3] - 从 VGG16 中得到 Generator 的输出图像特征[4] - 从 VGG16 中得到目标图像的特征[5] - 根据特征的差异性，计算损失函数[6] - BP将两张 Jeremy 和 Obama 图像对齐后，投影：tot_aligned_imgs = 2执行：python -W ignore stylegan2/epoching_custom_run_projector.py project-real-images --network=$model_path \ --dataset=custom_imgs --data-dir=datasets_stylegan2 --num-images=$tot_aligned_imgs --num-snapshots 500动态化过程的相关函数：def get_concat_h(im1, im2): dst = Image.new('RGB', (im1.width + im2.width, im1.height)) dst.paste(im1, (0, 0)) dst.paste(im2, (im1.width, 0)) return dst def make_project_progress_gifs(): all_result_folders = list(Path('results/').iterdir()) all_result_folders.sort() last_result_folder = all_result_folders[-1] for img_num in range(tot_aligned_imgs): all_step_pngs = [x for x in last_result_folder.iterdir() if x.name.endswith('png') and 'image{0:04d}'.format(img_num) in x.name] all_step_pngs.sort() target_image = Image.open(all_step_pngs[-1]).resize((results_size, results_size)) all_concat_imgs = [] for step_img_path in all_step_pngs[:-1]: step_img = Image.open(step_img_path).resize((results_size, results_size)) all_concat_imgs.append(get_concat_h(target_image, step_img)) # out_gif = 'output_gifs/image{0:04d}_project_progress.gif'.format(img_num) all_concat_imgs[0].save(out_gif, save_all=True, append_images=all_concat_imgs[1:], duration=1000/fps, loop=0) # make_project_progress_gifs()如：看下优化的 latent codes：def get_final_latents(): all_results = list(Path('results/').iterdir()) all_results.sort() last_result = all_results[-1] latent_files = [x for x in last_result.iterdir() if 'final_latent_code' in x.name] latent_files.sort() all_final_latents = [] for file in latent_files: with open(file, mode='rb') as latent_pickle: all_final_latents.append(pickle.load(latent_pickle)) return all_final_latents # latent_codes = get_final_latents() print(len(latent_codes), latent_codes[0].shape, latent_codes[1].shape) #(2, (1, 18, 512), (1, 18, 512))继续，看看从 latent codes 生成的图像是什么样.def load_model(): _G, _D, Gs = pretrained_networks.load_networks(model_path) noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')] Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) Gs_kwargs.randomize_noise = False return Gs, noise_vars, Gs_kwargs def generate_image_from_projected_latents(latent_vector): images = Gs.components.synthesis.run(latent_vector, **Gs_kwargs) return images # Gs, noise_vars, Gs_kwargs = load_model() # output_gifs_path = Path('output_gifs/') aligned_imgs_path = Path('aligned_imgs/') aligned_img_set = list(aligned_imgs_path.iterdir()) aligned_img_set.sort() aligned_img_set = [Image.open(x) for x in aligned_img_set] # images = generate_image_from_projected_latents(latent_codes[0]) recreated_img1 = Image.fromarray(images[0]).resize((results_size, results_size)) orig_img1 = aligned_img_set[1].resize((results_size, results_size)) get_concat_h(orig_img1, recreated_img1)如：images = generate_image_from_projected_latents(latent_codes[1]) recreated_img2 = Image.fromarray(images[0]).resize((results_size, results_size)) orig_img2 = aligned_img_set[0].resize((results_size, results_size)) get_concat_h(orig_img2, recreated_img2)如：重新运行下插值动画过程，def linear_interpolate(code1, code2, alpha): return code1 * alpha + code2 * (1 - alpha) def make_latent_interp_animation_real_faces(code1, code2, img1, img2, num_interps): step_size = 1.0/num_interps all_imgs = [] amounts = np.arange(0, 1, step_size) for alpha in tqdm(amounts): interpolated_latent_code = linear_interpolate(code1, code2, alpha) images = generate_image_from_projected_latents(interpolated_latent_code) interp_latent_image = Image.fromarray(images[0]).resize((400, 400)) frame = get_concat_h(img2, interp_latent_image) frame = get_concat_h(frame, img1) all_imgs.append(frame) save_name = output_gifs_path/'projected_latent_space_traversal.gif' all_imgs[0].save(save_name, save_all=True, append_images=all_imgs[1:], duration=1000/fps, loop=0) # make_latent_interp_animation_real_faces(latent_codes[0], latent_codes[1], recreated_img1, recreated_img2, num_interps=200)如：6. Latent Directions/Controls 修改投影图像已经有学习 latent space 中 latent directions 的方式，（有监督、无监督的都有），以控制特征.已经有开源的 StyleGAN2 directional latent vectors，可以通过在 latent space 中移动并控制某个特定特征.[1] - latent directions 的有监督学习首先，从模型中收集很多的样本（image + latent），并根据目标属性（如 smiling VS not smiling） 进行手工分类图片，并尝试保持样本表示的均衡性.然后，训练模型，分类或回归 latents 和手工标签. 这样，就可以采用学习的模型函数作为过渡方向（trasnition direction）.From: https://towardsdatascience.com/stylegan-v2-notes-on-training-and-latent-space-exploration-e51cf96584b3[2] - latent directions 的无监督学习Unsupervised Discovery of Interpretable Directions in the GAN Latent Space为了在 latent direction 进行移动，进行如下操作：latent_code = latent_code + latent_direction * magnitude # latent_code: # latent_direction: directional 向量，(18, 512)；其表示在 latent space 中移动的方向，并不是移动的大小 #magnitude: 在 latent_direction 方向移动的多少相比于两个点之间的插值，可以在特定的 latent direction 缓慢移动进行插值.相比于两个 latent codes 的融合，缓慢在 base latent code 上添加更多的 magnitude，并观察其变化.（这里采用 Robert Luxemburg 分享的 latent directions）def get_control_latent_vectors(path): files = [x for x in Path(path).iterdir() if str(x).endswith('.npy')] latent_vectors = {f.name[:-4]:np.load(f) for f in files} return latent_vectors # latent_controls = get_control_latent_vectors('stylegan2directions/') print(len(latent_controls), latent_controls.keys(), latent_controls['age'].shape) ''' (16, dict_keys(['age', 'eye_distance', 'eye_eyebrow_distance', 'eye_ratio', 'eyes_open', 'gender', 'lip_ratio', 'mouth_open', 'mouth_ratio', 'nose_mouth_distance', 'nose_ratio', 'nose_tip', 'pitch', 'roll', 'smile', 'yaw']), (18, 512)) ''' # def make_latent_control_animation(feature, start_amount, end_amount, step_size, person): all_imgs = [] amounts = np.linspace(start_amount, end_amount, abs(end_amount-start_amount)/step_size) for amount_to_move in tqdm(amounts): modified_latent_code = np.array(latent_code_to_use) modified_latent_code += latent_controls[feature]*amount_to_move images = generate_image_from_projected_latents(modified_latent_code) latent_img = Image.fromarray(images[0]).resize((results_size, results_size)) all_imgs.append(get_concat_h(image_to_use, latent_img)) save_name = output_gifs_path/'{0}_{1}.gif'.format(person, feature) all_imgs[0].save(save_name, save_all=True, append_images=all_imgs[1:], duration=1000/fps, loop=0) #有 16 个 latent "controls" #如 age latent_code_to_use = latent_codes[1] image_to_use = recreated_img2 make_latent_control_animation(feature='age', start_amount=-10, end_amount=10, step_size=0.1, person='jeremy') #如 age latent_code_to_use = latent_codes[0] image_to_use = recreated_img1 make_latent_control_animation(feature='age', start_amount=-5, end_amount=5, step_size=0.1, person='obama')如：[1] - 年龄变化[2] - 性别变化[3] - 嘴巴闭合变化[4] - 微笑变化
CycleGAN 流程：[1] - 从域 A 中选择一张图片input_A，通过Generator A2B 变换为域 B 中的一张假图片Generated_B，计算分类器 Discrimination_B 的损失;[2] - 将生成的假图片通过Generator B2A 变换回域 A，得到 Cyclic_A，计算其与 input_A 之间的损失；[3] - 另一网络，执行相同的流程处理.1. CycleGAN#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- """ Pix2pix. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function import os import time from absl import app from absl import flags import tensorflow as tf import tensorflow_datasets as tfds FLAGS = flags.FLAGS flags.DEFINE_integer('buffer_size', 400, 'Shuffle buffer size') flags.DEFINE_integer('batch_size', 1, 'Batch Size') flags.DEFINE_integer('epochs', 1, 'Number of epochs') flags.DEFINE_string('path', None, 'Path to the data folder') flags.DEFINE_boolean('enable_function', True, 'Enable Function?') IMG_WIDTH = 256 IMG_HEIGHT = 256 AUTOTUNE = tf.data.experimental.AUTOTUNE # 根据可用的CPU动态设置并行调用的数量 def random_crop(image): cropped_image = tf.image.random_crop(image, size=[IMG_HEIGHT, IMG_WIDTH, 3]) return cropped_image def normalize(image): image = tf.cast(image, tf.float32) image = (image / 127.5) - 1 return image @tf.function def random_jitter(image): # 调整大小为 286 x 286 x 3 image = tf.image.resize(image, [286, 286], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) # 随机裁剪到 256 x 256 x 3 image = random_crop(image) # 随机镜像 image = tf.image.random_flip_left_right(image) return image def preprocess_image_train(image, label): image = random_jitter(image) image = normalize(image) return image def preprocess_image_test(image, label): image = normalize(image) return image def create_horse2zebra_dataset(buffer_size, batch_size): """ Creates a tf.data Dataset. Args: path_to_train_images: Path to train images folder. path_to_test_images: Path to test images folder. buffer_size: Shuffle buffer size. batch_size: Batch size Returns: train dataset, test dataset """ dataset, metadata = tfds.load('cycle_gan/horse2zebra', with_info=True, as_supervised=True) train_horses, train_zebras = dataset['trainA'], dataset['trainB'] test_horses, test_zebras = dataset['testA'], dataset['testB'] # train_horses = train_horses.map(preprocess_image_train, num_parallel_calls=AUTOTUNE).cache() train_horses = train_horses.shuffle(buffer_size) train_horses = train_horses.batch(batch_size) train_zebras = train_zebras.map(preprocess_image_train, num_parallel_calls=AUTOTUNE).cache() train_zebras = train_zebras.shuffle(buffer_size) train_zebras = train_zebras.batch(batch_size) test_horses = test_horses.map(preprocess_image_test, num_parallel_calls=AUTOTUNE).cache() test_horses = test_horses.shuffle(buffer_size) test_horses = test_horses.batch(batch_size) test_zebras = test_zebras.map(preprocess_image_test, num_parallel_calls=AUTOTUNE).cache() test_zebras = test_zebras.shuffle(buffer_size) test_zebras = test_zebras.batch(batch_size) # dataset_dict = {} dataset_dict['train_horses'] = train_horses dataset_dict['train_zebras'] = train_zebras dataset_dict['test_horses'] = test_horses dataset_dict['test_zebras'] = test_zebras return dataset_dict class InstanceNormalization(tf.keras.layers.Layer): """ Instance Normalization Layer (https://arxiv.org/abs/1607.08022). """ def __init__(self, epsilon=1e-5): super(InstanceNormalization, self).__init__() self.epsilon = epsilon def build(self, input_shape): self.scale = self.add_weight( name='scale', shape=input_shape[-1:], initializer=tf.random_normal_initializer(1., 0.02), trainable=True) self.offset = self.add_weight( name='offset', shape=input_shape[-1:], initializer='zeros', trainable=True) def call(self, x): mean, variance = tf.nn.moments(x, axes=[1, 2], keepdims=True) inv = tf.math.rsqrt(variance + self.epsilon) normalized = (x - mean) * inv return self.scale * normalized + self.offset def downsample(filters, size, norm_type='batchnorm', apply_norm=True): """ Downsamples an input. Conv2D => Batchnorm => LeakyRelu Args: filters: number of filters size: filter size norm_type: Normalization type; either 'batchnorm' or 'instancenorm'. apply_norm: If True, adds the batchnorm layer Returns: Downsample Sequential Model """ initializer = tf.random_normal_initializer(0., 0.02) result = tf.keras.Sequential() result.add( tf.keras.layers.Conv2D(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)) if apply_norm: if norm_type.lower() == 'batchnorm': result.add(tf.keras.layers.BatchNormalization()) elif norm_type.lower() == 'instancenorm': result.add(InstanceNormalization()) result.add(tf.keras.layers.LeakyReLU()) return result def upsample(filters, size, norm_type='batchnorm', apply_dropout=False): """ Upsamples an input. Conv2DTranspose => Batchnorm => Dropout => Relu Args: filters: number of filters size: filter size norm_type: Normalization type; either 'batchnorm' or 'instancenorm'. apply_dropout: If True, adds the dropout layer Returns: Upsample Sequential Model """ initializer = tf.random_normal_initializer(0., 0.02) result = tf.keras.Sequential() result.add( tf.keras.layers.Conv2DTranspose(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)) if norm_type.lower() == 'batchnorm': result.add(tf.keras.layers.BatchNormalization()) elif norm_type.lower() == 'instancenorm': result.add(InstanceNormalization()) if apply_dropout: result.add(tf.keras.layers.Dropout(0.5)) result.add(tf.keras.layers.ReLU()) return result def unet_generator(output_channels, norm_type='batchnorm'): """ Modified u-net generator model (https://arxiv.org/abs/1611.07004). Args: output_channels: Output channels norm_type: Type of normalization. Either 'batchnorm' or 'instancenorm'. Returns: Generator model """ down_stack = [ downsample(64, 4, norm_type, apply_norm=False), # (bs, 128, 128, 64) downsample(128, 4, norm_type), # (bs, 64, 64, 128) downsample(256, 4, norm_type), # (bs, 32, 32, 256) downsample(512, 4, norm_type), # (bs, 16, 16, 512) downsample(512, 4, norm_type), # (bs, 8, 8, 512) downsample(512, 4, norm_type), # (bs, 4, 4, 512) downsample(512, 4, norm_type), # (bs, 2, 2, 512) downsample(512, 4, norm_type), # (bs, 1, 1, 512) ] up_stack = [ upsample(512, 4, norm_type, apply_dropout=True), # (bs, 2, 2, 1024) upsample(512, 4, norm_type, apply_dropout=True), # (bs, 4, 4, 1024) upsample(512, 4, norm_type, apply_dropout=True), # (bs, 8, 8, 1024) upsample(512, 4, norm_type), # (bs, 16, 16, 1024) upsample(256, 4, norm_type), # (bs, 32, 32, 512) upsample(128, 4, norm_type), # (bs, 64, 64, 256) upsample(64, 4, norm_type), # (bs, 128, 128, 128) ] initializer = tf.random_normal_initializer(0., 0.02) last = tf.keras.layers.Conv2DTranspose( output_channels, 4, strides=2, padding='same', kernel_initializer=initializer, activation='tanh') # (bs, 256, 256, 3) concat = tf.keras.layers.Concatenate() inputs = tf.keras.layers.Input(shape=[None, None, 3]) x = inputs # Downsampling through the model skips = [] for down in down_stack: x = down(x) skips.append(x) skips = reversed(skips[:-1]) # Upsampling and establishing the skip connections for up, skip in zip(up_stack, skips): x = up(x) x = concat([x, skip]) x = last(x) return tf.keras.Model(inputs=inputs, outputs=x) def discriminator(norm_type='batchnorm', target=True): """ PatchGan discriminator model (https://arxiv.org/abs/1611.07004). Args: norm_type: Type of normalization. Either 'batchnorm' or 'instancenorm'. target: Bool, indicating whether target image is an input or not. Returns: Discriminator model """ initializer = tf.random_normal_initializer(0., 0.02) inp = tf.keras.layers.Input(shape=[None, None, 3], name='input_image') x = inp if target: tar = tf.keras.layers.Input(shape=[None, None, 3], name='target_image') x = tf.keras.layers.concatenate([inp, tar]) # (bs, 256, 256, channels*2) down1 = downsample(64, 4, norm_type, False)(x) # (bs, 128, 128, 64) down2 = downsample(128, 4, norm_type)(down1) # (bs, 64, 64, 128) down3 = downsample(256, 4, norm_type)(down2) # (bs, 32, 32, 256) zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256) conv = tf.keras.layers.Conv2D( 512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1) # (bs, 31, 31, 512) if norm_type.lower() == 'batchnorm': norm1 = tf.keras.layers.BatchNormalization()(conv) elif norm_type.lower() == 'instancenorm': norm1 = InstanceNormalization()(conv) leaky_relu = tf.keras.layers.LeakyReLU()(norm1) zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512) last = tf.keras.layers.Conv2D( 1, 4, strides=1, kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1) if target: return tf.keras.Model(inputs=[inp, tar], outputs=last) else: return tf.keras.Model(inputs=inp, outputs=last) def get_checkpoint_prefix(): checkpoint_dir = './training_checkpoints' checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt') return checkpoint_prefix # class CycleGAN(object): def __init__(self, epochs, enable_function): self.epochs = epochs self.enable_function = enable_function self.lambda_value = 10 self.loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True) self.generator_g_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.generator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.discriminator_x_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.discriminator_y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.generator_g = unet_generator(output_channels=3, norm_type='instancenorm') self.generator_f = unet_generator(output_channels=3, norm_type='instancenorm') self.discriminator_x = discriminator(norm_type='instancenorm', target=False) self.discriminator_y = discriminator(norm_type='instancenorm', target=False) self.checkpoint = tf.train.Checkpoint( generator_g=self.generator_g, generator_f=self.generator_f, discriminator_x=self.discriminator_x, discriminator_y=self.discriminator_y, generator_g_optimizer=self.generator_g_optimizer, generator_f_optimizer=self.generator_f_optimizer, discriminator_x_optimizer=self.discriminator_x_optimizer, discriminator_y_optimizer=self.discriminator_y_optimizer) def discriminator_loss(self, disc_real_output, disc_generated_output): real_loss = self.loss_object(tf.ones_like(disc_real_output), disc_real_output) generated_loss = self.loss_object(tf.zeros_like(disc_generated_output), disc_generated_output) total_disc_loss = real_loss + generated_loss return total_disc_loss * 0.5 # def generator_loss(self, generated): gan_loss = self.loss_object(tf.ones_like(generated), generated) return gan_loss def calc_cycle_loss(self, real_image, cycled_image): loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image)) return self.lambda_value * loss1 def identity_loss(self, real_image, same_image): loss = tf.reduce_mean(tf.abs(real_image - same_image)) return self.lambda_value * 0.5 * loss @tf.function def train_step(self, real_x, real_y): # persistent 设置为 Ture，因为 GradientTape 被多次应用于计算梯度 with tf.GradientTape(persistent=True) as tape: # 生成器 G 转换 X -> Y # 生成器 F 转换 Y -> X fake_y = self.generator_g(real_x, training=True) cycled_x = self.generator_f(fake_y, training=True) fake_x = self.generator_f(real_y, training=True) cycled_y = self.generator_g(fake_x, training=True) # same_x 和 same_y 用于一致性损失。 same_x = self.generator_f(real_x, training=True) same_y = self.generator_g(real_y, training=True) disc_real_x = self.discriminator_x(real_x, training=True) disc_real_y = self.discriminator_y(real_y, training=True) disc_fake_x = self.discriminator_x(fake_x, training=True) disc_fake_y = self.discriminator_y(fake_y, training=True) # 计算损失 gen_g_loss = self.generator_loss(disc_fake_y) gen_f_loss = self.generator_loss(disc_fake_x) total_cycle_loss = self.calc_cycle_loss(real_x, cycled_x) + self.calc_cycle_loss(real_y, cycled_y) # 总生成器损失 = 对抗性损失 + 循环损失。 total_gen_g_loss = gen_g_loss + total_cycle_loss + self.identity_loss(real_y, same_y) total_gen_f_loss = gen_f_loss + total_cycle_loss + self.identity_loss(real_x, same_x) disc_x_loss = self.discriminator_loss(disc_real_x, disc_fake_x) disc_y_loss = self.discriminator_loss(disc_real_y, disc_fake_y) # 计算生成器和判别器损失。 generator_g_gradients = tape.gradient(total_gen_g_loss, self.generator_g.trainable_variables) generator_f_gradients = tape.gradient(total_gen_f_loss, self.generator_f.trainable_variables) discriminator_x_gradients = tape.gradient(disc_x_loss, self.discriminator_x.trainable_variables) discriminator_y_gradients = tape.gradient(disc_y_loss, self.discriminator_y.trainable_variables) # 将梯度应用于优化器。 self.generator_g_optimizer.apply_gradients(zip(generator_g_gradients, self.generator_g.trainable_variables)) self.generator_f_optimizer.apply_gradients(zip(generator_f_gradients, self.generator_f.trainable_variables)) self.discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients, self.discriminator_x.trainable_variables)) self.discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients, self.discriminator_y.trainable_variables)) return total_gen_g_loss, total_gen_f_loss, disc_x_loss, disc_y_loss def train(self, dataset, checkpoint_pr): """ Train the GAN for x number of epochs. Args: dataset: train dataset. checkpoint_pr: prefix in which the checkpoints are stored. Returns: Time for each epoch. """ time_list = [] if self.enable_function: self.train_step = tf.function(self.train_step) # train_horses, train_zebras = dataset['train_horses'], dataset['train_zebras'] for epoch in range(self.epochs): start_time = time.time() for image_x, image_y in tf.data.Dataset.zip((train_horses, train_zebras)): total_gen_g_loss, total_gen_f_loss, disc_x_loss, disc_y_loss = self.train_step(image_x, image_y) wall_time_sec = time.time() - start_time time_list.append(wall_time_sec) # saving (checkpoint) the model every 20 epochs if (epoch + 1) % 10 == 0: self.checkpoint.save(file_prefix=checkpoint_pr) template = 'Epoch {}, Generator G loss {}, Generator F loss {}, Discriminator X Loss {}, Discriminator Y Loss {}' print(template.format(epoch, total_gen_g_loss, total_gen_f_loss, disc_x_loss, disc_y_loss)) return time_list def run_main(argv): del argv kwargs = {'epochs': FLAGS.epochs, 'enable_function': FLAGS.enable_function, 'path': FLAGS.path, 'buffer_size': FLAGS.buffer_size, 'batch_size': FLAGS.batch_size} main(**kwargs) def main(epochs, enable_function, path, buffer_size, batch_size): path_to_folder = path cyclegan_object = CycleGAN(epochs, enable_function) dataset_dict = create_horse2zebra_dataset(buffer_size, batch_size) checkpoint_pr = get_checkpoint_prefix() print ('Training ...') return cyclegan_object.train(dataset_dict, checkpoint_pr) if __name__ == '__main__': app.run(run_main)
原文：TensorFlow - CycleGAN该文主要介绍使用CycleGAN 进行未配对图像到图像转换，参考论文 Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks - 2017， 提出了一种可以捕捉图像域特征并找出如何将这些特征转换为另一个图像域的方法，而无需任何成对的训练样本. 这里假定已经熟悉 Pix2Pix，可以参考 TF官方GAN教程-Pix2Pix. CycleGAN 的代码与其相似，主要区别在于额外的损失函数，以及非配对训练数据的使用. CycleGAN 使用循环一致损失来使训练过程无需配对数据. 换句话说，它可以从一个域转换到另一个域，而不需要在源域与目标域之间进行一对一映射. 这为完成许多有趣的任务开辟了可能性，例如照片增强、图片着色、风格迁移等. 所需要的只是源数据集和目标数据集(仅仅是图片目录).1. 配置pip install tensorflow_datasets #注 #pip install tensorflow-gpu==2.0 #pip install tensorflow_datasets==2.1.0导入相关库：import os import time import matplotlib.pyplot as plt import tensorflow as tf import tensorflow_datasets as tfds from tensorflow_examples.models.pix2pix import pix2pix tfds.disable_progress_bar() AUTOTUNE = tf.data.experimental.AUTOTUNE2. 输入 Pipline训练一个模型，以将普通马图片转换为斑马图片.如论文 Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks - 2017, 将随机抖动和镜像应用到训练集. 这是一些避免过拟合的图像增强技术. 类似于 TF官方GAN教程-Pix2Pix:[1] - 在随机抖动中，图片大小调整为 286 x 286，随后被随机裁剪为 256 x 256. [2] - 在随机镜像中，图片会从左到右随机翻转.dataset, metadata = tfds.load('cycle_gan/horse2zebra', with_info=True, as_supervised=True) train_horses, train_zebras = dataset['trainA'], dataset['trainB'] test_horses, test_zebras = dataset['testA'], dataset['testB'] # BUFFER_SIZE = 1000 BATCH_SIZE = 1 IMG_WIDTH = 256 IMG_HEIGHT = 256 # def random_crop(image): cropped_image = tf.image.random_crop( image, size=[IMG_HEIGHT, IMG_WIDTH, 3]) return cropped_image # 将图像归一化到区间 [-1, 1] 内. def normalize(image): image = tf.cast(image, tf.float32) image = (image / 127.5) - 1 return image def random_jitter(image): # 调整大小为 286 x 286 x 3 image = tf.image.resize(image, [286, 286], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) # 随机裁剪到 256 x 256 x 3 image = random_crop(image) # 随机镜像 image = tf.image.random_flip_left_right(image) return image def preprocess_image_train(image, label): image = random_jitter(image) image = normalize(image) return image def preprocess_image_test(image, label): image = normalize(image) return image # train_horses = train_horses.map( preprocess_image_train, num_parallel_calls=AUTOTUNE).cache().shuffle( BUFFER_SIZE).batch(1) train_zebras = train_zebras.map( preprocess_image_train, num_parallel_calls=AUTOTUNE).cache().shuffle( BUFFER_SIZE).batch(1) test_horses = test_horses.map( preprocess_image_test, num_parallel_calls=AUTOTUNE).cache().shuffle( BUFFER_SIZE).batch(1) test_zebras = test_zebras.map( preprocess_image_test, num_parallel_calls=AUTOTUNE).cache().shuffle( BUFFER_SIZE).batch(1) # sample_horse = next(iter(train_horses)) sample_zebra = next(iter(train_zebras)) plt.figure() plt.subplot(121) plt.title('Horse') plt.imshow(sample_horse[0] * 0.5 + 0.5) plt.subplot(122) plt.title('Horse with random jitter') plt.imshow(random_jitter(sample_horse[0]) * 0.5 + 0.5) plt.show() # plt.figure() plt.subplot(121) plt.title('Zebra') plt.imshow(sample_zebra[0] * 0.5 + 0.5) plt.subplot(122) plt.title('Zebra with random jitter') plt.imshow(random_jitter(sample_zebra[0]) * 0.5 + 0.5)3. 导入并重用 Pix2Pix 模型通过安装的 tensorflow_examples 包导入 Pix2Pix 中的生成器和判别器. 这里，使用模型体系结构与 pix2pix 中所使用的非常相似. 一些区别在于：[1] - Cyclegan 使用 instance normalization(实例归一化)而不是 batch normalization (批归一化). [2] - CycleGAN 论文使用一种基于 resnet 的改进生成器. 简单起见，这里使用的是改进的 unet 生成器. 这里训练了两个生成器(G 和 F)以及两个判别器(X 和 Y).生成器 G 学习将图片 X 转换为 Y. $G:(X \rightarrow Y)$生成器 F 学习将图片 Y 转换为 X. $F: Y \rightarrow X$判别器 D_X 学习区分图片 X 与生成的图片 $X(F(Y))$.判别器 D_Y 学习区分图片 Y 与生成的图片 $Y(G(X))$.OUTPUT_CHANNELS = 3 generator_g = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm') generator_f = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm') discriminator_x = pix2pix.discriminator(norm_type='instancenorm', target=False) discriminator_y = pix2pix.discriminator(norm_type='instancenorm', target=False) to_zebra = generator_g(sample_horse) to_horse = generator_f(sample_zebra) plt.figure(figsize=(8, 8)) contrast = 8 imgs = [sample_horse, to_zebra, sample_zebra, to_horse] title = ['Horse', 'To Zebra', 'Zebra', 'To Horse'] for i in range(len(imgs)): plt.subplot(2, 2, i+1) plt.title(title[i]) if i % 2 == 0: plt.imshow(imgs[i][0] * 0.5 + 0.5) else: plt.imshow(imgs[i][0] * 0.5 * contrast + 0.5) plt.show()plt.figure(figsize=(8, 8)) plt.subplot(121) plt.title('Is a real zebra?') plt.imshow(discriminator_y(sample_zebra)[0, ..., -1], cmap='RdBu_r') plt.subplot(122) plt.title('Is a real horse?') plt.imshow(discriminator_x(sample_horse)[0, ..., -1], cmap='RdBu_r') plt.show()4. 损失函数在 CycleGAN 中，没有可训练的成对数据，因此无法保证输入 x 和 目标 y 数据对在训练期间是有意义的. 所以为了强制网络学习正确的映射，作者提出了循环一致损失. 判别器损失和生成器损失和 pix2pix 中所使用的类似.LAMBDA = 10 loss_obj = tf.keras.losses.BinaryCrossentropy(from_logits=True) def discriminator_loss(real, generated): real_loss = loss_obj(tf.ones_like(real), real) generated_loss = loss_obj(tf.zeros_like(generated), generated) total_disc_loss = real_loss + generated_loss return total_disc_loss * 0.5 def generator_loss(generated): return loss_obj(tf.ones_like(generated), generated)循环一致意味着结果应接近原始输出. 例如，将一句英文译为法文，随后再从法文翻译回英文，最终的结果句应与原始句输入相同. 在循环一致损失中，[1] - 图片 X 通过生成器 G 传递，该生成器生成图片 $\hat{Y}$. [2] - 生成的图片 $\hat{Y}$ 通过生成器 F 传递，循环生成图片 $\hat{X}$. [3] - 在 X 和 $\hat{X}$ 之间计算平均绝对误差.$$ forward \ cycle \ consistency \ loss = X \rightarrow G(X) \rightarrow F(G(X)) \backsim \hat{X} $$$$ backward \ cycle \ consistency \ loss = Y \rightarrow F(Y) \rightarrow G(F(Y)) \backsim \hat{Y} $$def calc_cycle_loss(real_image, cycled_image): loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image)) return LAMBDA * loss1如上所示，生成器 G 负责将图片 X 转换为 Y. 一致性损失表明，如果将图片 Y 馈送给生成器 G，它应当生成真实图片 Y 或接近于 Y 的图片.$$ Identity \ loss = |G(Y) - Y| + |F(X) - X| $$def identity_loss(real_image, same_image): loss = tf.reduce_mean(tf.abs(real_image - same_image)) return LAMBDA * 0.5 * loss为所有生成器和判别器初始化优化器.generator_g_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) generator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) discriminator_x_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) discriminator_y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)5. Checkpointscheckpoint_path = "./checkpoints/train" ckpt = tf.train.Checkpoint(generator_g=generator_g, generator_f=generator_f, discriminator_x=discriminator_x, discriminator_y=discriminator_y, generator_g_optimizer=generator_g_optimizer, generator_f_optimizer=generator_f_optimizer, discriminator_x_optimizer=discriminator_x_optimizer, discriminator_y_optimizer=discriminator_y_optimizer) ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5) # 如果存在检查点，恢复最新版本检查点 if ckpt_manager.latest_checkpoint: ckpt.restore(ckpt_manager.latest_checkpoint) print ('Latest checkpoint restored!!')6. 训练这里仅对模型训练少量 epochs，如40, 预测的准确率可能会低一些. 而论文里为 200 epochs.EPOCHS = 40 def generate_images(model, test_input): prediction = model(test_input) plt.figure(figsize=(12, 12)) display_list = [test_input[0], prediction[0]] title = ['Input Image', 'Predicted Image'] for i in range(2): plt.subplot(1, 2, i+1) plt.title(title[i]) # 获取范围在 [0, 1] 之间的像素值以绘制它. plt.imshow(display_list[i] * 0.5 + 0.5) plt.axis('off') plt.show()尽管训练循环看起来很复杂，其实包含四个基本步骤：[1] - 获取预测. [2] - 计算损失值. [3] - 使用反向传播计算损失值. [4] - 将梯度应用于优化器.@tf.function def train_step(real_x, real_y): # persistent 设置为 Ture，因为 GradientTape 被多次应用于计算梯度. with tf.GradientTape(persistent=True) as tape: # 生成器 G 转换 X -> Y. # 生成器 F 转换 Y -> X. fake_y = generator_g(real_x, training=True) cycled_x = generator_f(fake_y, training=True) fake_x = generator_f(real_y, training=True) cycled_y = generator_g(fake_x, training=True) # same_x 和 same_y 用于一致性损失. same_x = generator_f(real_x, training=True) same_y = generator_g(real_y, training=True) disc_real_x = discriminator_x(real_x, training=True) disc_real_y = discriminator_y(real_y, training=True) disc_fake_x = discriminator_x(fake_x, training=True) disc_fake_y = discriminator_y(fake_y, training=True) # 计算损失. gen_g_loss = generator_loss(disc_fake_y) gen_f_loss = generator_loss(disc_fake_x) total_cycle_loss = calc_cycle_loss(real_x, cycled_x) + calc_cycle_loss(real_y, cycled_y) # 总生成器损失 = 对抗性损失 + 循环损失. total_gen_g_loss = gen_g_loss + total_cycle_loss + identity_loss(real_y, same_y) total_gen_f_loss = gen_f_loss + total_cycle_loss + identity_loss(real_x, same_x) disc_x_loss = discriminator_loss(disc_real_x, disc_fake_x) disc_y_loss = discriminator_loss(disc_real_y, disc_fake_y) # 计算生成器和判别器损失. generator_g_gradients = tape.gradient(total_gen_g_loss, generator_g.trainable_variables) generator_f_gradients = tape.gradient(total_gen_f_loss, generator_f.trainable_variables) discriminator_x_gradients = tape.gradient(disc_x_loss, discriminator_x.trainable_variables) discriminator_y_gradients = tape.gradient(disc_y_loss, discriminator_y.trainable_variables) # 将梯度应用于优化器. generator_g_optimizer.apply_gradients(zip(generator_g_gradients, generator_g.trainable_variables)) generator_f_optimizer.apply_gradients(zip(generator_f_gradients, generator_f.trainable_variables)) discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients, discriminator_x.trainable_variables)) discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients, discriminator_y.trainable_variables)) # for epoch in range(EPOCHS): start = time.time() n = 0 for image_x, image_y in tf.data.Dataset.zip((train_horses, train_zebras)): train_step(image_x, image_y) if n % 10 == 0: print ('.', end='') n+=1 # 使用一致的图像(sample_horse)，以便模型的进度清晰可见. generate_images(generator_g, sample_horse) if (epoch + 1) % 5 == 0: ckpt_save_path = ckpt_manager.save() print ('Saving checkpoint for epoch {} at {}'.format(epoch+1, ckpt_save_path)) print ('Time taken for epoch {} is {} sec\n'.format(epoch + 1, time.time()-start))7. 使用测试数据集进行生成# 在测试数据集上运行训练的模型. for inp in test_horses.take(5): generate_images(generator_g, inp)8. 总结这里简单基于 TF官方GAN教程-Pix2Pix 实现的生成器和判别器来实现 CycleGAN. 也可以尝试使用TensorFlow 数据集 中的不同的数据集. 也可以训练更多的 epoch 以改进结果，或者可以实现论文中所使用的改进 ResNet 生成器来代替这里使用的 U-Net 生成器.
出处：tensorflow_examples/models/pix2pix1. data_downloadhttps://github.com/tensorflow/examples/blob/master/tensorflow_examples/models/pix2pix/data_download.py""" Download facades data. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function import os from absl import app from absl import flags import tensorflow as tf FLAGS = flags.FLAGS flags.DEFINE_string('download_path', 'datasets', 'Download folder') _URL = 'https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz' def _main(argv): del argv download_path = FLAGS.download_path main(download_path) def main(download_path): path_to_zip = tf.keras.utils.get_file( 'facades.tar.gz', cache_subdir=download_path, origin=_URL, extract=True) path_to_folder = os.path.join(os.path.dirname(path_to_zip), 'facades/') return path_to_folder if __name__ == '__main__': app.run(_main)2. pix2pix#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- """ Pix2pix. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function import os import time from absl import app from absl import flags import tensorflow as tf FLAGS = flags.FLAGS flags.DEFINE_integer('buffer_size', 400, 'Shuffle buffer size') flags.DEFINE_integer('batch_size', 1, 'Batch Size') flags.DEFINE_integer('epochs', 1, 'Number of epochs') flags.DEFINE_string('path', None, 'Path to the data folder') flags.DEFINE_boolean('enable_function', True, 'Enable Function?') IMG_WIDTH = 256 IMG_HEIGHT = 256 AUTOTUNE = tf.data.experimental.AUTOTUNE #根据可用的CPU动态设置并行调用的数量 def load(image_file): """Loads the image and generates input and target image. Args: image_file: .jpeg file Returns: Input image, target image """ image = tf.io.read_file(image_file) image = tf.image.decode_jpeg(image) w = tf.shape(image)[1] w = w // 2 real_image = image[:, :w, :] input_image = image[:, w:, :] input_image = tf.cast(input_image, tf.float32) real_image = tf.cast(real_image, tf.float32) return input_image, real_image def resize(input_image, real_image, height, width): input_image = tf.image.resize(input_image, [height, width], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) real_image = tf.image.resize(real_image, [height, width], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) return input_image, real_image def random_crop(input_image, real_image): stacked_image = tf.stack([input_image, real_image], axis=0) cropped_image = tf.image.random_crop( stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3]) return cropped_image[0], cropped_image[1] def normalize(input_image, real_image): input_image = (input_image / 127.5) - 1 real_image = (real_image / 127.5) - 1 return input_image, real_image @tf.function def random_jitter(input_image, real_image): """ Random jittering. Resizes to 286 x 286 and then randomly crops to IMG_HEIGHT x IMG_WIDTH. Args: input_image: Input Image real_image: Real Image Returns: Input Image, real image """ # resizing to 286 x 286 x 3 input_image, real_image = resize(input_image, real_image, 286, 286) # randomly cropping to 256 x 256 x 3 input_image, real_image = random_crop(input_image, real_image) if tf.random.uniform(()) > 0.5: # random mirroring input_image = tf.image.flip_left_right(input_image) real_image = tf.image.flip_left_right(real_image) return input_image, real_image def load_image_train(image_file): input_image, real_image = load(image_file) input_image, real_image = random_jitter(input_image, real_image) input_image, real_image = normalize(input_image, real_image) return input_image, real_image def load_image_test(image_file): input_image, real_image = load(image_file) input_image, real_image = resize(input_image, real_image, IMG_HEIGHT, IMG_WIDTH) input_image, real_image = normalize(input_image, real_image) return input_image, real_image def create_dataset(path_to_train_images, path_to_test_images, buffer_size, batch_size): """Creates a tf.data Dataset. Args: path_to_train_images: Path to train images folder. path_to_test_images: Path to test images folder. buffer_size: Shuffle buffer size. batch_size: Batch size Returns: train dataset, test dataset """ train_dataset = tf.data.Dataset.list_files(path_to_train_images) train_dataset = train_dataset.shuffle(buffer_size) train_dataset = train_dataset.map( load_image_train, num_parallel_calls=AUTOTUNE) train_dataset = train_dataset.batch(batch_size) test_dataset = tf.data.Dataset.list_files(path_to_test_images) test_dataset = test_dataset.map( load_image_test, num_parallel_calls=AUTOTUNE) test_dataset = test_dataset.batch(batch_size) return train_dataset, test_dataset class InstanceNormalization(tf.keras.layers.Layer): """ Instance Normalization Layer (https://arxiv.org/abs/1607.08022). """ def __init__(self, epsilon=1e-5): super(InstanceNormalization, self).__init__() self.epsilon = epsilon def build(self, input_shape): self.scale = self.add_weight( name='scale', shape=input_shape[-1:], initializer=tf.random_normal_initializer(1., 0.02), trainable=True) self.offset = self.add_weight( name='offset', shape=input_shape[-1:], initializer='zeros', trainable=True) def call(self, x): mean, variance = tf.nn.moments(x, axes=[1, 2], keepdims=True) inv = tf.math.rsqrt(variance + self.epsilon) normalized = (x - mean) * inv return self.scale * normalized + self.offset def downsample(filters, size, norm_type='batchnorm', apply_norm=True): """ Downsamples an input. Conv2D => Batchnorm => LeakyRelu Args: filters: number of filters size: filter size norm_type: Normalization type; either 'batchnorm' or 'instancenorm'. apply_norm: If True, adds the batchnorm layer Returns: Downsample Sequential Model """ initializer = tf.random_normal_initializer(0., 0.02) result = tf.keras.Sequential() result.add( tf.keras.layers.Conv2D(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)) if apply_norm: if norm_type.lower() == 'batchnorm': result.add(tf.keras.layers.BatchNormalization()) elif norm_type.lower() == 'instancenorm': result.add(InstanceNormalization()) result.add(tf.keras.layers.LeakyReLU()) return result def upsample(filters, size, norm_type='batchnorm', apply_dropout=False): """ Upsamples an input. Conv2DTranspose => Batchnorm => Dropout => Relu Args: filters: number of filters size: filter size norm_type: Normalization type; either 'batchnorm' or 'instancenorm'. apply_dropout: If True, adds the dropout layer Returns: Upsample Sequential Model """ initializer = tf.random_normal_initializer(0., 0.02) result = tf.keras.Sequential() result.add( tf.keras.layers.Conv2DTranspose(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)) if norm_type.lower() == 'batchnorm': result.add(tf.keras.layers.BatchNormalization()) elif norm_type.lower() == 'instancenorm': result.add(InstanceNormalization()) if apply_dropout: result.add(tf.keras.layers.Dropout(0.5)) result.add(tf.keras.layers.ReLU()) return result def unet_generator(output_channels, norm_type='batchnorm'): """ Modified u-net generator model (https://arxiv.org/abs/1611.07004). Args: output_channels: Output channels norm_type: Type of normalization. Either 'batchnorm' or 'instancenorm'. Returns: Generator model """ down_stack = [ downsample(64, 4, norm_type, apply_norm=False), # (bs, 128, 128, 64) downsample(128, 4, norm_type), # (bs, 64, 64, 128) downsample(256, 4, norm_type), # (bs, 32, 32, 256) downsample(512, 4, norm_type), # (bs, 16, 16, 512) downsample(512, 4, norm_type), # (bs, 8, 8, 512) downsample(512, 4, norm_type), # (bs, 4, 4, 512) downsample(512, 4, norm_type), # (bs, 2, 2, 512) downsample(512, 4, norm_type), # (bs, 1, 1, 512) ] up_stack = [ upsample(512, 4, norm_type, apply_dropout=True), # (bs, 2, 2, 1024) upsample(512, 4, norm_type, apply_dropout=True), # (bs, 4, 4, 1024) upsample(512, 4, norm_type, apply_dropout=True), # (bs, 8, 8, 1024) upsample(512, 4, norm_type), # (bs, 16, 16, 1024) upsample(256, 4, norm_type), # (bs, 32, 32, 512) upsample(128, 4, norm_type), # (bs, 64, 64, 256) upsample(64, 4, norm_type), # (bs, 128, 128, 128) ] initializer = tf.random_normal_initializer(0., 0.02) last = tf.keras.layers.Conv2DTranspose( output_channels, 4, strides=2, padding='same', kernel_initializer=initializer, activation='tanh') # (bs, 256, 256, 3) concat = tf.keras.layers.Concatenate() inputs = tf.keras.layers.Input(shape=[None, None, 3]) x = inputs # Downsampling through the model skips = [] for down in down_stack: x = down(x) skips.append(x) skips = reversed(skips[:-1]) # Upsampling and establishing the skip connections for up, skip in zip(up_stack, skips): x = up(x) x = concat([x, skip]) x = last(x) return tf.keras.Model(inputs=inputs, outputs=x) def discriminator(norm_type='batchnorm', target=True): """ PatchGan discriminator model (https://arxiv.org/abs/1611.07004). Args: norm_type: Type of normalization. Either 'batchnorm' or 'instancenorm'. target: Bool, indicating whether target image is an input or not. Returns: Discriminator model """ initializer = tf.random_normal_initializer(0., 0.02) inp = tf.keras.layers.Input(shape=[None, None, 3], name='input_image') x = inp if target: tar = tf.keras.layers.Input(shape=[None, None, 3], name='target_image') x = tf.keras.layers.concatenate([inp, tar]) # (bs, 256, 256, channels*2) down1 = downsample(64, 4, norm_type, False)(x) # (bs, 128, 128, 64) down2 = downsample(128, 4, norm_type)(down1) # (bs, 64, 64, 128) down3 = downsample(256, 4, norm_type)(down2) # (bs, 32, 32, 256) zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256) conv = tf.keras.layers.Conv2D( 512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1) # (bs, 31, 31, 512) if norm_type.lower() == 'batchnorm': norm1 = tf.keras.layers.BatchNormalization()(conv) elif norm_type.lower() == 'instancenorm': norm1 = InstanceNormalization()(conv) leaky_relu = tf.keras.layers.LeakyReLU()(norm1) zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512) last = tf.keras.layers.Conv2D( 1, 4, strides=1, kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1) if target: return tf.keras.Model(inputs=[inp, tar], outputs=last) else: return tf.keras.Model(inputs=inp, outputs=last) def get_checkpoint_prefix(): checkpoint_dir = './training_checkpoints' checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt') return checkpoint_prefix class Pix2pix(object): """ Pix2pix class. Args: epochs: Number of epochs. enable_function: If true, train step is decorated with tf.function. buffer_size: Shuffle buffer size.. batch_size: Batch size. """ def __init__(self, epochs, enable_function): self.epochs = epochs self.enable_function = enable_function self.lambda_value = 100 self.loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True) self.generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.generator = unet_generator(output_channels=3) self.discriminator = discriminator() self.checkpoint = tf.train.Checkpoint( generator_optimizer=self.generator_optimizer, discriminator_optimizer=self.discriminator_optimizer, generator=self.generator, discriminator=self.discriminator) def discriminator_loss(self, disc_real_output, disc_generated_output): real_loss = self.loss_object(tf.ones_like(disc_real_output), disc_real_output) generated_loss = self.loss_object(tf.zeros_like(disc_generated_output), disc_generated_output) total_disc_loss = real_loss + generated_loss return total_disc_loss def generator_loss(self, disc_generated_output, gen_output, target): gan_loss = self.loss_object(tf.ones_like(disc_generated_output), disc_generated_output) # mean absolute error l1_loss = tf.reduce_mean(tf.abs(target - gen_output)) total_gen_loss = gan_loss + (self.lambda_value * l1_loss) return total_gen_loss def train_step(self, input_image, target_image): """ One train step over the generator and discriminator model. Args: input_image: Input Image. target_image: Target image. Returns: generator loss, discriminator loss. """ with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape: gen_output = self.generator(input_image, training=True) disc_real_output = self.discriminator( [input_image, target_image], training=True) disc_generated_output = self.discriminator( [input_image, gen_output], training=True) gen_loss = self.generator_loss( disc_generated_output, gen_output, target_image) disc_loss = self.discriminator_loss( disc_real_output, disc_generated_output) generator_gradients = gen_tape.gradient( gen_loss, self.generator.trainable_variables) discriminator_gradients = disc_tape.gradient( disc_loss, self.discriminator.trainable_variables) self.generator_optimizer.apply_gradients(zip( generator_gradients, self.generator.trainable_variables)) self.discriminator_optimizer.apply_gradients(zip( discriminator_gradients, self.discriminator.trainable_variables)) return gen_loss, disc_loss def train(self, dataset, checkpoint_pr): """ Train the GAN for x number of epochs. Args: dataset: train dataset. checkpoint_pr: prefix in which the checkpoints are stored. Returns: Time for each epoch. """ time_list = [] if self.enable_function: self.train_step = tf.function(self.train_step) for epoch in range(self.epochs): start_time = time.time() for input_image, target_image in dataset: gen_loss, disc_loss = self.train_step(input_image, target_image) wall_time_sec = time.time() - start_time time_list.append(wall_time_sec) # saving (checkpoint) the model every 20 epochs if (epoch + 1) % 20 == 0: self.checkpoint.save(file_prefix=checkpoint_pr) template = 'Epoch {}, Generator loss {}, Discriminator Loss {}' print (template.format(epoch, gen_loss, disc_loss)) return time_list def run_main(argv): del argv kwargs = {'epochs': FLAGS.epochs, 'enable_function': FLAGS.enable_function, 'path': FLAGS.path, 'buffer_size': FLAGS.buffer_size, 'batch_size': FLAGS.batch_size} main(**kwargs) def main(epochs, enable_function, path, buffer_size, batch_size): path_to_folder = path pix2pix_object = Pix2pix(epochs, enable_function) train_dataset, _ = create_dataset( os.path.join(path_to_folder, 'train/*.jpg'), os.path.join(path_to_folder, 'test/*.jpg'), buffer_size, batch_size) checkpoint_pr = get_checkpoint_prefix() print ('Training ...') return pix2pix_object.train(train_dataset, checkpoint_pr) if __name__ == '__main__': app.run(run_main)3. pix2pix_test""" Tests for Pix2Pix. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function from absl import flags import tensorflow as tf from pix2pix import data_download from pix2pix import pix2pix FLAGS = flags.FLAGS class Pix2PixTest(tf.test.TestCase): def test_one_step_with_function(self): epochs = 1 batch_size = 1 enable_function = True input_image = tf.random.uniform((256, 256, 3)) target_image = tf.random.uniform((256, 256, 3)) train_dataset = tf.data.Dataset.from_tensors( (input_image, target_image)).map(pix2pix.random_jitter).batch(batch_size) checkpoint_pr = pix2pix.get_checkpoint_prefix() pix2pix_obj = pix2pix.Pix2pix(epochs, enable_function) pix2pix_obj.train(train_dataset, checkpoint_pr) def test_one_step_without_function(self): epochs = 1 batch_size = 1 enable_function = False input_image = tf.random.uniform((256, 256, 3)) target_image = tf.random.uniform((256, 256, 3)) train_dataset = tf.data.Dataset.from_tensors( (input_image, target_image)).map(pix2pix.random_jitter).batch(batch_size) pix2pix_obj = pix2pix.Pix2pix(epochs, enable_function) checkpoint_pr = pix2pix.get_checkpoint_prefix() pix2pix_obj.train(train_dataset, checkpoint_pr) class Pix2PixBenchmark(tf.test.Benchmark): def __init__(self, output_dir=None, **kwargs): self.output_dir = output_dir def benchmark_with_function(self): path = data_download.main("datasets") kwargs = {"epochs": 6, "enable_function": True, "path": path, "buffer_size": 400, "batch_size": 1} self._run_and_report_benchmark(**kwargs) def benchmark_without_function(self): path = data_download.main("datasets") kwargs = {"epochs": 6, "enable_function": False, "path": path, "buffer_size": 400, "batch_size": 1} self._run_and_report_benchmark(**kwargs) def _run_and_report_benchmark(self, **kwargs): time_list = pix2pix.main(**kwargs) # 1st epoch is the warmup epoch hence skipping it for calculating time. self.report_benchmark(wall_time=tf.reduce_mean(time_list[1:])) if __name__ == "__main__": tf.test.main()
原文: TensorFlow - Pix2Pix该文主要是介绍采用 conditional GAN 进行图像转换(image to image translation), 可参考 Image-to-Image Translation with Conditional Adversarial Networks - 2016基于该技术,可以实现黑白图像的彩色化, 将谷歌地图转换为谷歌地球,等等. 这里,以建筑立面(building facades) 转换为真实建筑图(real buildings).采用的数据集为:CMP Facade Database 的 Image-to-Image Translation with Conditional Adversarial Networks - 2016 论文作者提供的一个预处理版本: https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/.每个 epoch 在 V100 GPU 上大概耗时 15s.如下图, 为训练 200 epochs 后模型生成的例图:1. 配置import os import time import matplotlib.pyplot as plt import tensorflow as tf2. 加载数据集2.1. 数据增强对于训练数据集, 进行随机抖动(random jittering) 和镜像(mirroring)处理.[1] - 随机抖动, 图像先缩放为286x286, 再随机裁剪 256x256.[2] - 随机镜像, 图像随机水平翻转,如, 左到右._URL = 'https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz' path_to_zip = tf.keras.utils.get_file('facades.tar.gz', origin=_URL, extract=True) PATH = os.path.join(os.path.dirname(path_to_zip), 'facades/') # BUFFER_SIZE = 400 BATCH_SIZE = 1 IMG_WIDTH = 256 IMG_HEIGHT = 256 # def load(image_file): image = tf.io.read_file(image_file) image = tf.image.decode_jpeg(image) w = tf.shape(image)[1] w = w // 2 real_image = image[:, :w, :] input_image = image[:, w:, :] input_image = tf.cast(input_image, tf.float32) real_image = tf.cast(real_image, tf.float32) return input_image, real_image # inp, re = load(PATH+'train/100.jpg') # casting to int for matplotlib to show the image plt.figure() plt.imshow(inp/255.0) plt.figure() plt.imshow(re/255.0)def resize(input_image, real_image, height, width): input_image = tf.image.resize(input_image, [height, width], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) real_image = tf.image.resize(real_image, [height, width], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) return input_image, real_image def random_crop(input_image, real_image): stacked_image = tf.stack([input_image, real_image], axis=0) cropped_image = tf.image.random_crop( stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3]) return cropped_image[0], cropped_image[1] # normalizing the images to [-1, 1] def normalize(input_image, real_image): input_image = (input_image / 127.5) - 1 real_image = (real_image / 127.5) - 1 return input_image, real_image @tf.function() def random_jitter(input_image, real_image): # resizing to 286 x 286 x 3 input_image, real_image = resize(input_image, real_image, 286, 286) # randomly cropping to 256 x 256 x 3 input_image, real_image = random_crop(input_image, real_image) if tf.random.uniform(()) > 0.5: # random mirroring input_image = tf.image.flip_left_right(input_image) real_image = tf.image.flip_left_right(real_image) return input_image, real_image如图, 图像的处理流程为:[1] - 将图像尺寸调整到更大的 height 和 width；[2] - 随机裁剪到目标尺寸；[3] - 随机水平翻转图像.plt.figure(figsize=(6, 6)) for i in range(4): rj_inp, rj_re = random_jitter(inp, re) plt.subplot(2, 2, i+1) plt.imshow(rj_inp/255.0) plt.axis('off') plt.show()2.2. 数据集加载def load_image_train(image_file): input_image, real_image = load(image_file) input_image, real_image = random_jitter(input_image, real_image) input_image, real_image = normalize(input_image, real_image) return input_image, real_image def load_image_test(image_file): input_image, real_image = load(image_file) input_image, real_image = resize(input_image, real_image, IMG_HEIGHT, IMG_WIDTH) input_image, real_image = normalize(input_image, real_image) return input_image, real_image3. 输入 Pipelines# train dataset train_dataset = tf.data.Dataset.list_files(PATH+'train/*.jpg') train_dataset = train_dataset.map( load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE) train_dataset = train_dataset.shuffle(BUFFER_SIZE) train_dataset = train_dataset.batch(BATCH_SIZE) # test dataset test_dataset = tf.data.Dataset.list_files(PATH+'test/*.jpg') test_dataset = test_dataset.map(load_image_test) test_dataset = test_dataset.batch(BATCH_SIZE)4. 构建生成器[1] - 网络结构是 U-Net 的改良.[2] - 编码器(encoder)的每个 block 结构为：Conv -> Batchnorm -> Leaky ReLU.[3] - 解码器(decoder)的每个 block 结构为：Transposed Conv -> Batchnorm -> Dropout(applied to the first 3 blocks) -> ReLU.[4] - 编码器和解码器之间采用跳跃连接(skip connection)，类似于 U-Net.OUTPUT_CHANNELS = 3 #编码器 def downsample(filters, size, apply_batchnorm=True): initializer = tf.random_normal_initializer(0., 0.02) result = tf.keras.Sequential() result.add( tf.keras.layers.Conv2D(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)) if apply_batchnorm: result.add(tf.keras.layers.BatchNormalization()) result.add(tf.keras.layers.LeakyReLU()) return result # down_model = downsample(3, 4) down_result = down_model(tf.expand_dims(inp, 0)) print (down_result.shape) # (1, 128, 128, 3) #解码器 def upsample(filters, size, apply_dropout=False): initializer = tf.random_normal_initializer(0., 0.02) result = tf.keras.Sequential() result.add( tf.keras.layers.Conv2DTranspose(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)) result.add(tf.keras.layers.BatchNormalization()) if apply_dropout: result.add(tf.keras.layers.Dropout(0.5)) result.add(tf.keras.layers.ReLU()) return result # up_model = upsample(3, 4) up_result = up_model(down_result) print (up_result.shape) # (1, 256, 256, 3) #生成器 def Generator(): inputs = tf.keras.layers.Input(shape=[256,256,3]) down_stack = [ downsample(64, 4, apply_batchnorm=False), # (bs, 128, 128, 64) downsample(128, 4), # (bs, 64, 64, 128) downsample(256, 4), # (bs, 32, 32, 256) downsample(512, 4), # (bs, 16, 16, 512) downsample(512, 4), # (bs, 8, 8, 512) downsample(512, 4), # (bs, 4, 4, 512) downsample(512, 4), # (bs, 2, 2, 512) downsample(512, 4), # (bs, 1, 1, 512) ] up_stack = [ upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024) upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024) upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024) upsample(512, 4), # (bs, 16, 16, 1024) upsample(256, 4), # (bs, 32, 32, 512) upsample(128, 4), # (bs, 64, 64, 256) upsample(64, 4), # (bs, 128, 128, 128) ] initializer = tf.random_normal_initializer(0., 0.02) last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4, strides=2, padding='same', kernel_initializer=initializer, activation='tanh') # (bs, 256, 256, 3) x = inputs # Downsampling through the model skips = [] for down in down_stack: x = down(x) skips.append(x) skips = reversed(skips[:-1]) # Upsampling and establishing the skip connections for up, skip in zip(up_stack, skips): x = up(x) x = tf.keras.layers.Concatenate()([x, skip]) x = last(x) return tf.keras.Model(inputs=inputs, outputs=x) # generator = Generator() tf.keras.utils.plot_model(generator, show_shapes=True, dpi=64) # gen_output = generator(inp[tf.newaxis,...], training=False) plt.imshow(gen_output[0,...])4.1. 生成器损失函数[1] - 计算的是生成图像和全是1数组的 sigmoid 交叉熵损失函数.[2] - Image-to-Image Translation with Conditional Adversarial Networks - 2016 还采用了 L1 损失函数，计算生成图像和目标图像间的 MAE(均方误差).[3] - 生成器总损失函数的计算为：$$ loss = gan\_loss + \lambda * l1\_loss $$其中，$\lambda=100$.生成器损失函数实现如：LAMBDA = 100 def generator_loss(disc_generated_output, gen_output, target): gan_loss = loss_object(tf.ones_like(disc_generated_output), disc_generated_output) # mean absolute error l1_loss = tf.reduce_mean(tf.abs(target - gen_output)) total_gen_loss = gan_loss + (LAMBDA * l1_loss) return total_gen_loss, gan_loss, l1_loss流程如下：5. 构建判别器[1] - 判别器是 PatchGAN;[2] - 判别器每个 block 结构为：Conv -> BatchNorm -> Leaky ReLU[3] - 判别器最后一层的输出 shape 为：(batch_size, 30, 30, 1)[4] - 输出的每个 30x30 patch 分别对输入图像的 70x70 部分进行分类(这种结构被称为 PatchGAN)[5] - 判别器包含两个输入：输入图像和目标图像，其应该被分类为 real.输入图像和生成图像，其应该被分类为 fake.采用 tf.concat([inp, tar], axis=-1) 将两个输入连接.def Discriminator(): initializer = tf.random_normal_initializer(0., 0.02) inp = tf.keras.layers.Input(shape=[256, 256, 3], name='input_image') tar = tf.keras.layers.Input(shape=[256, 256, 3], name='target_image') x = tf.keras.layers.concatenate([inp, tar]) # (bs, 256, 256, channels*2) down1 = downsample(64, 4, False)(x) # (bs, 128, 128, 64) down2 = downsample(128, 4)(down1) # (bs, 64, 64, 128) down3 = downsample(256, 4)(down2) # (bs, 32, 32, 256) zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256) conv = tf.keras.layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1) # (bs, 31, 31, 512) batchnorm1 = tf.keras.layers.BatchNormalization()(conv) leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1) zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512) last = tf.keras.layers.Conv2D(1, 4, strides=1, kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1) return tf.keras.Model(inputs=[inp, tar], outputs=last) # discriminator = Discriminator() tf.keras.utils.plot_model(discriminator, show_shapes=True, dpi=64) # disc_out = discriminator([inp[tf.newaxis,...], gen_output], training=False) plt.imshow(disc_out[0,...,-1], vmin=-20, vmax=20, cmap='RdBu_r') plt.colorbar()5.1. 判别器损失函数[1] - 判别器损失函数两个输入： real 图像和生成图像.[2] - real_loss 为 real 图像与全1数组之间的 sigmoid 交叉熵损失函数.[3] - generated_loss 为生成图像与全0数组之间的 sigmoid 交叉熵损失函数.[4] - total_loss 为 real_loss 和 generated_loss 的之和.loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True) def discriminator_loss(disc_real_output, disc_generated_output): real_loss = loss_object(tf.ones_like(disc_real_output), disc_real_output) generated_loss = loss_object(tf.zeros_like(disc_generated_output), disc_generated_output) total_disc_loss = real_loss + generated_loss return total_disc_loss流程如下：6. 定义 Optimizers 和 Checkpoint-savergenerator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) checkpoint_dir = './training_checkpoints' checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt") checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer, discriminator_optimizer=discriminator_optimizer, generator=generator, discriminator=discriminator)7. 生成图像创建一个训练过程中图像可视化的函数:[1] - 将测试数据集图像送到生成器；[2] - 生成器将其输入图像转化为输出；[3] - 最后可视化预测图像.def generate_images(model, test_input, tar): prediction = model(test_input, training=True) plt.figure(figsize=(15,15)) display_list = [test_input[0], tar[0], prediction[0]] title = ['Input Image', 'Ground Truth', 'Predicted Image'] for i in range(3): plt.subplot(1, 3, i+1) plt.title(title[i]) # getting the pixel values between [0, 1] to plot it. plt.imshow(display_list[i] * 0.5 + 0.5) plt.axis('off') plt.show() # for example_input, example_target in test_dataset.take(1): generate_images(generator, example_input, example_target)8. 训练[1] - 每个输入样本生成一个输出；[2] - 输入图像和生成的图像作为判别器的第一个输入；输入图像和目标图像作为判别器的第二个输入；[3] - 然后，计算生成器损失函数和判别器损失函数；[4] - 接着，计算关于生成器和判别器变量的梯度，并应用于优化器(optimizer)；[5] - TensorBoard 记录损失函数日志.实现如下：EPOCHS = 150 import datetime log_dir="logs/" summary_writer = tf.summary.create_file_writer( log_dir + "fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")) # @tf.function def train_step(input_image, target, epoch): with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape: gen_output = generator(input_image, training=True) disc_real_output = discriminator([input_image, target], training=True) disc_generated_output = discriminator([input_image, gen_output], training=True) gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(disc_generated_output, gen_output, target) disc_loss = discriminator_loss(disc_real_output, disc_generated_output) generator_gradients = gen_tape.gradient(gen_total_loss, generator.trainable_variables) discriminator_gradients = disc_tape.gradient(disc_loss, discriminator.trainable_variables) generator_optimizer.apply_gradients(zip(generator_gradients, generator.trainable_variables)) discriminator_optimizer.apply_gradients(zip(discriminator_gradients, discriminator.trainable_variables)) with summary_writer.as_default(): tf.summary.scalar('gen_total_loss', gen_total_loss, step=epoch) tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=epoch) tf.summary.scalar('gen_l1_loss', gen_l1_loss, step=epoch) tf.summary.scalar('disc_loss', disc_loss, step=epoch)8.1. 训练主循环[1] - epochs 迭代[2] - 每个 epoch，清除显示的图片，并运行generate_images 以显示训练进度；[3] - 每个 epoch，迭代处理训练数据集，并打印..[4] - 每 20 epochs 保存一次 checkpoint.def fit(train_ds, epochs, test_ds): for epoch in range(epochs): start = time.time() display.clear_output(wait=True) for example_input, example_target in test_ds.take(1): generate_images(generator, example_input, example_target) print("Epoch: ", epoch) # Train for n, (input_image, target) in train_ds.enumerate(): print('.', end='') if (n+1) % 100 == 0: print() train_step(input_image, target, epoch) print() # saving (checkpoint) the model every 20 epochs if (epoch + 1) % 20 == 0: checkpoint.save(file_prefix = checkpoint_prefix) print ('Time taken for epoch {} is {} sec\n'.format(epoch + 1, time.time()-start)) checkpoint.save(file_prefix = checkpoint_prefix) # fit(train_dataset, EPOCHS, test_dataset) # # restoring the latest checkpoint in checkpoint_dir checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))训练循环会保存日志信息，可以在 TensorBoard 监视训练过程. 9. 测试数据集生成图像# Run the trained model on a few examples from the test dataset for inp, tar in test_dataset.take(5): generate_images(generator, inp, tar)10. 完整实现#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- import os import time import matplotlib.pyplot as plt import datetime import tensorflow as tf print(tf.__version__) #2.3.0 # def load(image_file): image = tf.io.read_file(image_file) image = tf.image.decode_jpeg(image) w = tf.shape(image)[1] w = w // 2 real_image = image[:, :w, :] input_image = image[:, w:, :] input_image = tf.cast(input_image, tf.float32) real_image = tf.cast(real_image, tf.float32) return input_image, real_image def resize(input_image, real_image, height, width): input_image = tf.image.resize(input_image, [height, width], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) real_image = tf.image.resize(real_image, [height, width], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) return input_image, real_image def random_crop(input_image, real_image, IMG_HEIGHT = 256, IMG_WIDTH = 256): stacked_image = tf.stack([input_image, real_image], axis=0) cropped_image = tf.image.random_crop(stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3]) return cropped_image[0], cropped_image[1] # normalizing the images to [-1, 1] def normalize(input_image, real_image): input_image = (input_image / 127.5) - 1 real_image = (real_image / 127.5) - 1 return input_image, real_image @tf.function() def random_jitter(input_image, real_image): # resizing to 286 x 286 x 3 input_image, real_image = resize(input_image, real_image, 286, 286) # randomly cropping to 256 x 256 x 3 input_image, real_image = random_crop(input_image, real_image) if tf.random.uniform(()) > 0.5: # random mirroring input_image = tf.image.flip_left_right(input_image) real_image = tf.image.flip_left_right(real_image) return input_image, real_image # def load_image_train(image_file): input_image, real_image = load(image_file) input_image, real_image = random_jitter(input_image, real_image) input_image, real_image = normalize(input_image, real_image) return input_image, real_image def load_image_test(image_file, IMG_HEIGHT = 256, IMG_WIDTH = 256): input_image, real_image = load(image_file) input_image, real_image = resize(input_image, real_image, IMG_HEIGHT, IMG_WIDTH) input_image, real_image = normalize(input_image, real_image) return input_image, real_image #编码器 def downsample(filters, size, apply_batchnorm=True): initializer = tf.random_normal_initializer(0., 0.02) result = tf.keras.Sequential() result.add(tf.keras.layers.Conv2D(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)) if apply_batchnorm: result.add(tf.keras.layers.BatchNormalization()) result.add(tf.keras.layers.LeakyReLU()) return result #解码器 def upsample(filters, size, apply_dropout=False): initializer = tf.random_normal_initializer(0., 0.02) result = tf.keras.Sequential() result.add(tf.keras.layers.Conv2DTranspose(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)) result.add(tf.keras.layers.BatchNormalization()) if apply_dropout: result.add(tf.keras.layers.Dropout(0.5)) result.add(tf.keras.layers.ReLU()) return result #生成器 def Generator(OUTPUT_CHANNELS = 3): inputs = tf.keras.layers.Input(shape=[256, 256, 3]) down_stack = [ downsample(64, 4, apply_batchnorm=False), # (bs, 128, 128, 64) downsample(128, 4), # (bs, 64, 64, 128) downsample(256, 4), # (bs, 32, 32, 256) downsample(512, 4), # (bs, 16, 16, 512) downsample(512, 4), # (bs, 8, 8, 512) downsample(512, 4), # (bs, 4, 4, 512) downsample(512, 4), # (bs, 2, 2, 512) downsample(512, 4), # (bs, 1, 1, 512) ] up_stack = [ upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024) upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024) upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024) upsample(512, 4), # (bs, 16, 16, 1024) upsample(256, 4), # (bs, 32, 32, 512) upsample(128, 4), # (bs, 64, 64, 256) upsample(64, 4), # (bs, 128, 128, 128) ] initializer = tf.random_normal_initializer(0., 0.02) last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4, strides=2, padding='same', kernel_initializer=initializer, activation='tanh') # (bs, 256, 256, 3) x = inputs # Downsampling through the model skips = [] for down in down_stack: x = down(x) skips.append(x) skips = reversed(skips[:-1]) # Upsampling and establishing the skip connections for up, skip in zip(up_stack, skips): x = up(x) x = tf.keras.layers.Concatenate()([x, skip]) x = last(x) return tf.keras.Model(inputs=inputs, outputs=x) # def Discriminator(): initializer = tf.random_normal_initializer(0., 0.02) inp = tf.keras.layers.Input(shape=[256, 256, 3], name='input_image') tar = tf.keras.layers.Input(shape=[256, 256, 3], name='target_image') x = tf.keras.layers.concatenate([inp, tar]) # (bs, 256, 256, channels*2) down1 = downsample(64, 4, False)(x) # (bs, 128, 128, 64) down2 = downsample(128, 4)(down1) # (bs, 64, 64, 128) down3 = downsample(256, 4)(down2) # (bs, 32, 32, 256) zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256) conv = tf.keras.layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1) # (bs, 31, 31, 512) batchnorm1 = tf.keras.layers.BatchNormalization()(conv) leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1) zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512) last = tf.keras.layers.Conv2D(1, 4, strides=1, kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1) return tf.keras.Model(inputs=[inp, tar], outputs=last) # loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True) def generator_loss(disc_generated_output, gen_output, target, LAMBDA = 100): gan_loss = loss_object(tf.ones_like(disc_generated_output), disc_generated_output) # mean absolute error l1_loss = tf.reduce_mean(tf.abs(target - gen_output)) total_gen_loss = gan_loss + (LAMBDA * l1_loss) return total_gen_loss, gan_loss, l1_loss def discriminator_loss(disc_real_output, disc_generated_output): real_loss = loss_object(tf.ones_like(disc_real_output), disc_real_output) generated_loss = loss_object(tf.zeros_like(disc_generated_output), disc_generated_output) total_disc_loss = real_loss + generated_loss return total_disc_loss # def generate_images(model, test_input, tar): prediction = model(test_input, training=True) plt.figure(figsize=(15, 15)) display_list = [test_input[0], tar[0], prediction[0]] title = ['Input Image', 'Ground Truth', 'Predicted Image'] for i in range(3): plt.subplot(1, 3, i + 1) plt.title(title[i]) # getting the pixel values between [0, 1] to plot it. plt.imshow(display_list[i] * 0.5 + 0.5) plt.axis('off') # plt.show() # @tf.function def train_step(input_image, target, epoch): with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape: gen_output = generator(input_image, training=True) disc_real_output = discriminator([input_image, target], training=True) disc_generated_output = discriminator([input_image, gen_output], training=True) gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(disc_generated_output, gen_output, target) disc_loss = discriminator_loss(disc_real_output, disc_generated_output) generator_gradients = gen_tape.gradient(gen_total_loss, generator.trainable_variables) discriminator_gradients = disc_tape.gradient(disc_loss, discriminator.trainable_variables) generator_optimizer.apply_gradients(zip(generator_gradients, generator.trainable_variables)) discriminator_optimizer.apply_gradients(zip(discriminator_gradients, discriminator.trainable_variables)) with summary_writer.as_default(): tf.summary.scalar('gen_total_loss', gen_total_loss, step=epoch) tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=epoch) tf.summary.scalar('gen_l1_loss', gen_l1_loss, step=epoch) tf.summary.scalar('disc_loss', disc_loss, step=epoch) # def fit(train_ds, epochs, test_ds): for epoch in range(epochs): start = time.time() for example_input, example_target in test_ds.take(1): generate_images(generator, example_input, example_target) plt.savefig('epoch_{}.png'.format(epoch), format='png',) print("Epoch: ", epoch) # Train for n, (input_image, target) in train_ds.enumerate(): print('.', end='') if (n + 1) % 100 == 0: print() train_step(input_image, target, epoch) print() # saving (checkpoint) the model every 20 epochs if (epoch + 1) % 20 == 0: checkpoint.save(file_prefix=checkpoint_prefix) print('Time taken for epoch {} is {} sec\n'.format(epoch + 1, time.time() - start)) checkpoint.save(file_prefix=checkpoint_prefix) # _URL = 'https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz' path_to_zip = tf.keras.utils.get_file('facades.tar.gz', origin=_URL, extract=True) PATH = os.path.join(os.path.dirname(path_to_zip), 'facades/') # inp, re = load(PATH+'train/100.jpg') # # casting to int for matplotlib to show the image # plt.figure() # plt.imshow(inp/255.0) # plt.figure() # plt.imshow(re/255.0) # # # # plt.figure(figsize=(6, 6)) # for i in range(4): # rj_inp, rj_re = random_jitter(inp, re) # plt.subplot(2, 2, i+1) # plt.imshow(rj_inp/255.0) # plt.axis('off') # plt.show() # BATCH_SIZE = 1 # train dataset train_dataset = tf.data.Dataset.list_files(PATH+'train/*.jpg') train_dataset = train_dataset.map(load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE) train_dataset = train_dataset.shuffle(400) #BUFFER_SIZE = 400 train_dataset = train_dataset.batch(BATCH_SIZE) # test dataset test_dataset = tf.data.Dataset.list_files(PATH+'test/*.jpg') test_dataset = test_dataset.map(load_image_test) test_dataset = test_dataset.batch(BATCH_SIZE) # # # down_model = downsample(3, 4) # down_result = down_model(tf.expand_dims(inp, 0)) # print (down_result.shape) # # (1, 128, 128, 3) # # # up_model = upsample(3, 4) # up_result = up_model(down_result) # print (up_result.shape) # # (1, 256, 256, 3) # generator = Generator() # tf.keras.utils.plot_model(generator, show_shapes=True, dpi=64) # discriminator = Discriminator() # tf.keras.utils.plot_model(discriminator, show_shapes=True, dpi=64) # # # gen_output = generator(inp[tf.newaxis,...], training=False) # plt.imshow(gen_output[0,...]) # # # # disc_out = discriminator([inp[tf.newaxis,...], gen_output], training=False) # plt.imshow(disc_out[0,...,-1], vmin=-20, vmax=20, cmap='RdBu_r') # plt.colorbar() # generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) checkpoint_dir = './training_checkpoints' checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt") checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer, discriminator_optimizer=discriminator_optimizer, generator=generator, discriminator=discriminator) # # # for example_input, example_target in test_dataset.take(1): # generate_images(generator, example_input, example_target) # log_dir="logs/" summary_writer = tf.summary.create_file_writer( log_dir + "fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")) # EPOCHS = 150 fit(train_dataset, EPOCHS, test_dataset) # # restoring the latest checkpoint in checkpoint_dir checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir)) # Run the trained model on a few examples from the test dataset for inp, tar in test_dataset.take(5): generate_images(generator, inp, tar)
From: tensorflow_examples/models/dcgan1. dcgan.py#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- """ DCGAN. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function import os import time from absl import app from absl import flags import tensorflow as tf import tensorflow_datasets as tfds FLAGS = flags.FLAGS flags.DEFINE_integer('buffer_size', 10000, 'Shuffle buffer size') flags.DEFINE_integer('batch_size', 64, 'Batch Size') flags.DEFINE_integer('epochs', 1, 'Number of epochs') flags.DEFINE_boolean('enable_function', True, 'Enable Function?') AUTOTUNE = tf.data.experimental.AUTOTUNE def scale(image, label): image = tf.cast(image, tf.float32) image = (image - 127.5) / 127.5 return image, label def create_dataset(buffer_size, batch_size): train_dataset = tfds.load('mnist', split='train', as_supervised=True, shuffle_files=True) train_dataset = train_dataset.map(scale, num_parallel_calls=AUTOTUNE) train_dataset = train_dataset.shuffle(buffer_size).batch(batch_size) return train_dataset def make_generator_model(): """ Generator. Returns: Keras Sequential model """ model = tf.keras.Sequential([ tf.keras.layers.Dense(7*7*256, use_bias=False), tf.keras.layers.BatchNormalization(), tf.keras.layers.LeakyReLU(), tf.keras.layers.Reshape((7, 7, 256)), tf.keras.layers.Conv2DTranspose(128, 5, strides=(1, 1), padding='same', use_bias=False), tf.keras.layers.BatchNormalization(), tf.keras.layers.LeakyReLU(), tf.keras.layers.Conv2DTranspose(64, 5, strides=(2, 2), padding='same', use_bias=False), tf.keras.layers.BatchNormalization(), tf.keras.layers.LeakyReLU(), tf.keras.layers.Conv2DTranspose(1, 5, strides=(2, 2), padding='same', use_bias=False, activation='tanh') ]) return model def make_discriminator_model(): """ Discriminator. Returns: Keras Sequential model """ model = tf.keras.Sequential([ tf.keras.layers.Conv2D(64, 5, strides=(2, 2), padding='same'), tf.keras.layers.LeakyReLU(), tf.keras.layers.Dropout(0.3), tf.keras.layers.Conv2D(128, 5, strides=(2, 2), padding='same'), tf.keras.layers.LeakyReLU(), tf.keras.layers.Dropout(0.3), tf.keras.layers.Flatten(), tf.keras.layers.Dense(1) ]) return model def get_checkpoint_prefix(): checkpoint_dir = './training_checkpoints' checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt') return checkpoint_prefix class Dcgan(object): """ Dcgan class. Args: epochs: Number of epochs. enable_function: If true, train step is decorated with tf.function. batch_size: Batch size. """ def __init__(self, epochs, enable_function, batch_size): self.epochs = epochs self.enable_function = enable_function self.batch_size = batch_size self.noise_dim = 100 self.loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True) self.generator_optimizer = tf.keras.optimizers.Adam(1e-4) self.discriminator_optimizer = tf.keras.optimizers.Adam(1e-4) self.generator = make_generator_model() self.discriminator = make_discriminator_model() self.checkpoint = tf.train.Checkpoint( generator_optimizer=self.generator_optimizer, discriminator_optimizer=self.discriminator_optimizer, generator=self.generator, discriminator=self.discriminator) def generator_loss(self, generated_output): return self.loss_object(tf.ones_like(generated_output), generated_output) def discriminator_loss(self, real_output, generated_output): real_loss = self.loss_object(tf.ones_like(real_output), real_output) generated_loss = self.loss_object(tf.zeros_like(generated_output), generated_output) total_loss = real_loss + generated_loss return total_loss def train_step(self, image): """ One train step over the generator and discriminator model. Args: image: Input image. Returns: generator loss, discriminator loss. """ noise = tf.random.normal([self.batch_size, self.noise_dim]) with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape: generated_images = self.generator(noise, training=True) real_output = self.discriminator(image, training=True) generated_output = self.discriminator(generated_images, training=True) gen_loss = self.generator_loss(generated_output) disc_loss = self.discriminator_loss(real_output, generated_output) gradients_of_generator = gen_tape.gradient( gen_loss, self.generator.trainable_variables) gradients_of_discriminator = disc_tape.gradient( disc_loss, self.discriminator.trainable_variables) self.generator_optimizer.apply_gradients(zip( gradients_of_generator, self.generator.trainable_variables)) self.discriminator_optimizer.apply_gradients(zip( gradients_of_discriminator, self.discriminator.trainable_variables)) return gen_loss, disc_loss def train(self, dataset, checkpoint_pr): """ Train the GAN for x number of epochs. Args: dataset: train dataset. checkpoint_pr: prefix in which the checkpoints are stored. Returns: Time for each epoch. """ time_list = [] if self.enable_function: self.train_step = tf.function(self.train_step) for epoch in range(self.epochs): start_time = time.time() for image, _ in dataset: gen_loss, disc_loss = self.train_step(image) wall_time_sec = time.time() - start_time time_list.append(wall_time_sec) # saving (checkpoint) the model every 15 epochs if (epoch + 1) % 15 == 0: self.checkpoint.save(file_prefix=checkpoint_pr) template = 'Epoch {}, Generator loss {}, Discriminator Loss {}' print (template.format(epoch, gen_loss, disc_loss)) return time_list def run_main(argv): del argv kwargs = {'epochs': FLAGS.epochs, 'enable_function': FLAGS.enable_function, 'buffer_size': FLAGS.buffer_size, 'batch_size': FLAGS.batch_size} main(**kwargs) def main(epochs, enable_function, buffer_size, batch_size): train_dataset = create_dataset(buffer_size, batch_size) checkpoint_pr = get_checkpoint_prefix() dcgan_obj = Dcgan(epochs, enable_function, batch_size) print ('Training ...') return dcgan_obj.train(train_dataset, checkpoint_pr) if __name__ == '__main__': app.run(run_main)2. dcgan_test.py""" DCGAN tests. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function from absl import flags import tensorflow as tf from tensorflow_examples.models.dcgan import dcgan FLAGS = flags.FLAGS class DcganTest(tf.test.TestCase): def test_one_epoch_with_function(self): epochs = 1 batch_size = 1 enable_function = True input_image = tf.random.uniform((28, 28, 1)) label = tf.zeros((1,)) train_dataset = tf.data.Dataset.from_tensors( (input_image, label)).batch(batch_size) checkpoint_pr = dcgan.get_checkpoint_prefix() dcgan_obj = dcgan.Dcgan(epochs, enable_function, batch_size) dcgan_obj.train(train_dataset, checkpoint_pr) def test_one_epoch_without_function(self): epochs = 1 batch_size = 1 enable_function = False input_image = tf.random.uniform((28, 28, 1)) label = tf.zeros((1,)) train_dataset = tf.data.Dataset.from_tensors( (input_image, label)).batch(batch_size) checkpoint_pr = dcgan.get_checkpoint_prefix() dcgan_obj = dcgan.Dcgan(epochs, enable_function, batch_size) dcgan_obj.train(train_dataset, checkpoint_pr) class DCGANBenchmark(tf.test.Benchmark): def __init__(self, output_dir=None, **kwargs): self.output_dir = output_dir def benchmark_with_function(self): kwargs = {"epochs": 6, "enable_function": True, "buffer_size": 10000, "batch_size": 64} self._run_and_report_benchmark(**kwargs) def benchmark_without_function(self): kwargs = {"epochs": 6, "enable_function": False, "buffer_size": 10000, "batch_size": 64} self._run_and_report_benchmark(**kwargs) def _run_and_report_benchmark(self, **kwargs): time_list = dcgan.main(**kwargs) # 1st epoch is the warmup epoch hence skipping it for calculating time. self.report_benchmark(wall_time=tf.reduce_mean(time_list[1:])) if __name__ == "__main__": tf.test.main()
原文：TensorFlow - 深度卷积生成对抗网络该文主要是介绍如何使用深度卷积生成对抗网络(DCGAN)生成手写数字图片. 代码是使用 Keras Sequential API 与 tf.GradientTape 训练循环编写的.1. 什么是生成对抗网络生成对抗网络(GANs)是当今计算机科学领域最有趣的想法之一. 两个模型通过对抗过程同时训练. 一个生成器(“艺术家”)学习创造看起来真实的图像，而判别器(“艺术评论家”)学习区分真假图像. 训练过程中，生成器在生成逼真图像方面逐渐变强，而判别器在辨别这些图像的能力上逐渐变强. 当判别器不再能够区分真实图片和伪造图片时，训练过程达到平衡. 本文在 MNIST 数据集上演示了该过程. 下方动画展示了当训练了 50 个epoch (全部数据集迭代50次) 时生成器所生成的一系列图片. 图片从随机噪声开始，随着时间的推移越来越像手写数字. 关于 GANs 更多信息：参阅 MIT的 深度学习入门 课程.2. 配置首先安装用于生成 GIF 图片的库 - imageio:pip install -q imageioimport glob import imageio import matplotlib.pyplot as plt import numpy as np import os import PIL import time import tensorflow as tf print(tf.__version__) #2.3.0 from tensorflow.keras import layers3. 加载和准备数据集# MNIST 数据集 (train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data() train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32') train_images = (train_images - 127.5) / 127.5 # 将图片标准化到 [-1, 1] 区间内 BUFFER_SIZE = 60000 BATCH_SIZE = 256 # 批量化和打乱数据 train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)4. 创建模型生成器和判别器均使用 Keras Sequential API 定义.4.1. 生成器生成器使用 tf.keras.layers.Conv2DTranspose (上采样)层来从种子(随机噪声)中产生图片. 以一个使用该种子作为输入的 Dense 层开始，然后多次上采样直到达到所期望的 28x28x1 的图片尺寸. 注：除了输出层使用 tanh 之外，其他每层均使用 tf.keras.layers.LeakyReLU 作为激活函数.def make_generator_model(): model = tf.keras.Sequential() model.add(layers.Dense(7*7*256, use_bias=False, input_shape=(100,))) model.add(layers.BatchNormalization()) model.add(layers.LeakyReLU()) model.add(layers.Reshape((7, 7, 256))) assert model.output_shape == (None, 7, 7, 256) # 注意：batch size 没有限制 model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False)) assert model.output_shape == (None, 7, 7, 128) model.add(layers.BatchNormalization()) model.add(layers.LeakyReLU()) model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False)) assert model.output_shape == (None, 14, 14, 64) model.add(layers.BatchNormalization()) model.add(layers.LeakyReLU()) model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh')) assert model.output_shape == (None, 28, 28, 1) return model使用(尚未训练的)生成器创建一张图片.generator = make_generator_model() noise = tf.random.normal([1, 100]) generated_image = generator(noise, training=False) plt.imshow(generated_image[0, :, :, 0], cmap='gray') plt.show()如图：4.2. 判别器判别器是一个基于 CNN 的图片分类器.def make_discriminator_model(): model = tf.keras.Sequential() model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[28, 28, 1])) model.add(layers.LeakyReLU()) model.add(layers.Dropout(0.3)) model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same')) model.add(layers.LeakyReLU()) model.add(layers.Dropout(0.3)) model.add(layers.Flatten()) model.add(layers.Dense(1)) return model使用(尚未训练的)判别器来对图片的真伪进行判断. 模型将被训练为为真实图片输出正值，为伪造图片输出负值.discriminator = make_discriminator_model() decision = discriminator(generated_image) print (decision) #tf.Tensor([[-0.00427552]], shape=(1, 1), dtype=float32)5. 定义损失函数和优化器为两个模型定义损失函数和优化器.# 计算交叉熵损失的辅助函数 cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)5.1. 判别器损失判别器损失量化判别器从判断真伪图片的能力. 它将判别器对真实图片的预测值与值全为 1 的数组进行对比，将判别器对伪造(生成的)图片的预测值与值全为 0 的数组进行对比.def discriminator_loss(real_output, fake_output): real_loss = cross_entropy(tf.ones_like(real_output), real_output) fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output) total_loss = real_loss + fake_loss return total_loss5.2. 生成器损失生成器损失量化其欺骗判别器的能力. 直观来讲，如果生成器表现良好，判别器将会把伪造图片判断为真实图片(或 1). 这里将把判别器在生成图片上的判断结果与一个值全为 1 的数组进行对比.def generator_loss(fake_output): return cross_entropy(tf.ones_like(fake_output), fake_output)由于需要分别训练两个网络，判别器和生成器的优化器是不同的.generator_optimizer = tf.keras.optimizers.Adam(1e-4) discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)5.3. 保存checkpoint保存和恢复模型，在长时间训练任务被中断的情况下比较有帮助.checkpoint_dir = './training_checkpoints' checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt") checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer, discriminator_optimizer=discriminator_optimizer, generator=generator, discriminator=discriminator)6. 定义训练循环EPOCHS = 50 noise_dim = 100 num_examples_to_generate = 16 # 将重复使用该种子(因此在动画 GIF 中更容易可视化进度) seed = tf.random.normal([num_examples_to_generate, noise_dim])训练循环在生成器接收到一个随机种子作为输入时开始. 该种子用于生产一张图片. 判别器随后被用于区分真实图片(选自训练集)和伪造图片(由生成器生成). 针对这里的每一个模型都计算损失函数，并且计算梯度用于更新生成器与判别器.# 注意 `tf.function` 的使用 # 该注解使函数被“编译” @tf.function def train_step(images): noise = tf.random.normal([BATCH_SIZE, noise_dim]) with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape: generated_images = generator(noise, training=True) real_output = discriminator(images, training=True) fake_output = discriminator(generated_images, training=True) gen_loss = generator_loss(fake_output) disc_loss = discriminator_loss(real_output, fake_output) gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables) gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables) generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables)) discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables)) # def generate_and_save_images(model, epoch, test_input): # 注意 training 设定为 False # 因此，所有层都在推理模式下运行(batchnorm). predictions = model(test_input, training=False) fig = plt.figure(figsize=(4,4)) for i in range(predictions.shape[0]): plt.subplot(4, 4, i+1) plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray') plt.axis('off') plt.savefig('image_at_epoch_{:04d}.png'.format(epoch)) plt.show() # def train(dataset, epochs): for epoch in range(epochs): start = time.time() for image_batch in dataset: train_step(image_batch) # 继续进行时为 GIF 生成图像 generate_and_save_images(generator, epoch + 1, seed) # 每 15 个 epoch 保存一次模型 if (epoch + 1) % 15 == 0: checkpoint.save(file_prefix = checkpoint_prefix) print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start)) # 最后一个 epoch 结束后生成图片 generate_and_save_images(generator, epochs, seed)7. 训练模型调用上面定义的 train() 方法来同时训练生成器和判别器. 注：训练 GANs 可能是棘手的. 重要的是，生成器和判别器不能够互相压制对方(例如，他们以相似的学习率训练). 在训练之初，生成的图片看起来像是随机噪声. 随着训练过程的进行，生成的数字将越来越真实. 在大概 50 个 epoch 之后，这些图片看起来像是 MNIST 数字. 使用 Colab 中的默认设置可能需要大约 1 分钟每 epoch.train(train_dataset, EPOCHS)恢复最新的checkpoint.checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))9. 完整实现#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- import glob import imageio import matplotlib.pyplot as plt import numpy as np import os import PIL import time import tensorflow as tf print(tf.__version__) #2.3.0 from tensorflow.keras import layers # 生成网络 def make_generator_model(): model = tf.keras.Sequential() model.add(layers.Dense(7 * 7 * 256, use_bias=False, input_shape=(100,))) model.add(layers.BatchNormalization()) model.add(layers.LeakyReLU()) model.add(layers.Reshape((7, 7, 256))) assert model.output_shape == (None, 7, 7, 256) # 注意：batch size 没有限制 model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False)) assert model.output_shape == (None, 7, 7, 128) model.add(layers.BatchNormalization()) model.add(layers.LeakyReLU()) model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False)) assert model.output_shape == (None, 14, 14, 64) model.add(layers.BatchNormalization()) model.add(layers.LeakyReLU()) model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh')) assert model.output_shape == (None, 28, 28, 1) return model # 判别网络 def make_discriminator_model(): model = tf.keras.Sequential() model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[28, 28, 1])) model.add(layers.LeakyReLU()) model.add(layers.Dropout(0.3)) model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same')) model.add(layers.LeakyReLU()) model.add(layers.Dropout(0.3)) model.add(layers.Flatten()) model.add(layers.Dense(1)) return model # # 计算交叉熵损失的辅助函数 cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True) # 生成器损失函数 def generator_loss(fake_output): return cross_entropy(tf.ones_like(fake_output), fake_output) # 判别器损失函数 def discriminator_loss(real_output, fake_output): real_loss = cross_entropy(tf.ones_like(real_output), real_output) fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output) total_loss = real_loss + fake_loss return total_loss # def generate_and_save_images(model, epoch, test_input): # 注意 training 设定为 False # 因此，所有层都在推理模式下运行(batchnorm). predictions = model(test_input, training=False) plt.figure(figsize=(4, 4)) for i in range(predictions.shape[0]): plt.subplot(4, 4, i + 1) plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray') plt.axis('off') plt.savefig('image_at_epoch_{:04d}.png'.format(epoch)) plt.show() # @tf.function def train_step(images, generator, discriminator, generator_optimizer, discriminator_optimizer, batch_size, noise_dim): noise = tf.random.normal([batch_size, noise_dim]) with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape: generated_images = generator(noise, training=True) real_output = discriminator(images, training=True) fake_output = discriminator(generated_images, training=True) gen_loss = generator_loss(fake_output) disc_loss = discriminator_loss(real_output, fake_output) gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables) gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables) generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables)) discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables)) # def main(): # MNIST 数据集 (train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data() train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32') train_images = (train_images - 127.5) / 127.5 # 将图片标准化到 [-1, 1] 区间内 BUFFER_SIZE = 60000 BATCH_SIZE = 256 # 批量化和打乱数据 train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE) # generator = make_generator_model() noise = tf.random.normal([1, 100]) generated_image = generator(noise, training=False) # # plt.imshow(generated_image[0, :, :, 0], cmap='gray') # plt.show() # discriminator = make_discriminator_model() # decision = discriminator(generated_image) # print(decision) # generator_optimizer = tf.keras.optimizers.Adam(1e-4) discriminator_optimizer = tf.keras.optimizers.Adam(1e-4) # checkpoint_dir = './training_checkpoints' checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt") checkpoint = tf.train.Checkpoint( generator_optimizer=generator_optimizer, discriminator_optimizer=discriminator_optimizer, generator=generator, discriminator=discriminator) # EPOCHS = 50 noise_dim = 100 num_examples_to_generate = 16 # 将重复使用该种子 seed = tf.random.normal([num_examples_to_generate, noise_dim]) # train for epoch in range(EPOCHS): start = time.time() for image_batch in train_dataset: train_step(image_batch, generator, discriminator, generator_optimizer, discriminator_optimizer, BATCH_SIZE, noise_dim) # 继续进行时为 GIF 生成图像 generate_and_save_images(generator, epoch + 1, seed) # 每 15 个 epoch 保存一次模型 if (epoch + 1) % 15 == 0: checkpoint.save(file_prefix=checkpoint_prefix) print('Time for epoch {} is {} sec'.format(epoch + 1, time.time() - start)) # 最后一个 epoch 结束后生成图片 generate_and_save_images(generator, epochs, seed) #checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir)) print('[INFO]Done.') # if __name__ == '__main__': main()9. 创建 GIF# 使用 epoch 数生成单张图片 def display_image(epoch_no): return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no)) # display_image(EPOCHS)使用训练过程中生成的图片通过 imageio 生成动态 gif:anim_file = 'dcgan.gif' with imageio.get_writer(anim_file, mode='I') as writer: filenames = glob.glob('image*.png') filenames = sorted(filenames) last = -1 for i,filename in enumerate(filenames): frame = 2*(i**0.5) if round(frame) > round(last): last = frame else: continue image = imageio.imread(filename) writer.append_data(image) image = imageio.imread(filename) writer.append_data(image)10. 相关材料[1] - 大规模名人面部属性(CelebA)数据集 - Kaggle[2] - NIPS 2016 教程： 生成对抗网络
原文：TensorFlow - DeepDream该文是关于 DeepDream 的简单实现，可参考 Alexander Mordvintsev 的 blog - Inceptionism: Going Deeper into Neural Networks.DeepDream 是对于神经网络所学习的模式(patterns)的可视化实验. 类似于孩子看着天空中的云，并尝试解释随机形状时，DeepDream 会过度解释并增强其在图像中看到的图案.DeepDream 首先将图像进行网络的前向计算，然后计算图像关于特定网络层的激活值的梯度. 则图像被修改，以增强这些激活值，增强网络可见的模式，并生成类似梦境的图像. 该过程被称为"Inceptionism".下面开始说明如何进行神经网络梦境化，并增强其在图像中看到的超现实模式.1. 配置import numpy as np import matplotlib.pyplot as plt import matplotlib as mpl from PIL import Image import tensorflow as tf from tensorflow.keras.preprocessing import image2. 选择待梦境化的图片url = 'https://storage.googleapis.com/download.tensorflow.org/example_images/YellowLabradorLooking_new.jpg' # Download an image and read it into a NumPy array. def download(url, max_dim=None): name = url.split('/')[-1] image_path = tf.keras.utils.get_file(name, origin=url) img = PIL.Image.open(image_path) if max_dim: img.thumbnail((max_dim, max_dim)) return np.array(img) # Normalize an image def deprocess(img): img = 255*(img + 1.0)/2.0 return tf.cast(img, tf.uint8) # Display an image def show(img): img = Image.fromarray(np.array(img)) plt.imshow(img) plt.show() # 图像下载 original_img = download(url, max_dim=500) show(original_img)3. 特征提取模型下载一个预训练的图像分类模型. 可以根据需要选择任意预训练模型， 仅需要调整网络层的名字.这里采用类似于 DeepDream 原始采用的 InceptionV3 模型.base_model = tf.keras.applications.InceptionV3(include_top=False, weights='imagenet')DeepDream 的思想是，选择一个网络层(或多个网络层)，以图像逐渐“刺激(excites)” 网络层的方式最大化损失函数. 合并的特征复杂度取决于所选择的网络层，如，越浅(lower)的网络层得到图像中简单的模式，而越深(deeper)的网络层得到图像中更复杂的特征.InceptionV3 网络结构是比较大的. 对于DeepDream 而言，其感兴趣的是连接卷积层的网络层. InceptionV3 中有 11 个这样的网络层，网络层名为 "mixed0" 到 "mixed10". 采用不同的网络层会得到不同的梦境化图像. 越深网络层对应于 higher-level 特征(例如，eyes 和 faces)，而越浅网络层对应于简单特征(如边缘、形状、纹理等). 其中，越深网络层的训练时间越久，因为梯度计算更深.# Maximize the activations of these layers names = ['mixed3', 'mixed5'] layers = [base_model.get_layer(name).output for name in names] # Create the feature extraction model dream_model = tf.keras.Model(inputs=base_model.input, outputs=layers)4. 计算损失损失值是所选择网络层的激活值之和. 损失值在每一层进行了归一化，因此，数值较大的网络层不会比数值较小的网络层对损失值的作用更大. 一般来说，损失值是通过梯度下降进行最小化的标准. 但 DeepDream 中，是通过梯度上升(Gradient ascent)对损失函数最大化.def calc_loss(img, model): # Pass forward the image through the model to retrieve the activations. # Converts the image into a batch of size 1. img_batch = tf.expand_dims(img, axis=0) layer_activations = model(img_batch) if len(layer_activations) == 1: layer_activations = [layer_activations] losses = [] for act in layer_activations: loss = tf.math.reduce_mean(act) losses.append(loss) return tf.reduce_sum(losses)5. 梯度上升当计算了所选择网络层的损失值后，然后就是计算关于图像的梯度，并将其与原始图像相加.添加梯度到图像上，可以增强网络所看到的模式. 在每一步，都会创建一张新图片，其递增的刺激网络中特定网络层的激活值.class DeepDream(tf.Module): def __init__(self, model): self.model = model #性能起见，进行 tf.function封装， @tf.function( input_signature=( tf.TensorSpec(shape=[None,None,3], dtype=tf.float32), tf.TensorSpec(shape=[], dtype=tf.int32), tf.TensorSpec(shape=[], dtype=tf.float32),) ) def __call__(self, img, steps, step_size): print("Tracing") loss = tf.constant(0.0) for n in tf.range(steps): with tf.GradientTape() as tape: # This needs gradients relative to `img` # `GradientTape` only watches `tf.Variable`s by default tape.watch(img) loss = calc_loss(img, self.model) # 计算关于输入图像像素的loss的梯度 gradients = tape.gradient(loss, img) # Normalize the gradients. gradients /= tf.math.reduce_std(gradients) + 1e-8 # In gradient ascent, the "loss" is maximized so that the input image increasingly "excites" the layers. # You can update the image by directly adding the gradients (because they're the same shape!) img = img + gradients*step_size img = tf.clip_by_value(img, -1, 1) return loss, img # deepdream = DeepDream(dream_model)Better performance with tf.function6. 主循环def run_deep_dream_simple(img, steps=100, step_size=0.01): # Convert from uint8 to the range expected by the model. img = tf.keras.applications.inception_v3.preprocess_input(img) img = tf.convert_to_tensor(img) step_size = tf.convert_to_tensor(step_size) steps_remaining = steps step = 0 while steps_remaining: if steps_remaining>100: run_steps = tf.constant(100) else: run_steps = tf.constant(steps_remaining) steps_remaining -= run_steps step += run_steps loss, img = deepdream(img, run_steps, tf.constant(step_size)) show(deprocess(img)) print ("Step {}, loss {}".format(step, loss)) # result = deprocess(img) show(result) return result # dream_img = run_deep_dream_simple(img=original_img, steps=100, step_size=0.01)7. 升级版前面的实现存在一些问题：[1] - 输出是存在噪声的(其可以通过添加 tf.image.total_variation 损失来解决)[2] - 图像是低分辨率的[3] - 模式看起来好像都是发生在相同的粒度上.解决这些问题的一个方法是，在不同的尺度进行梯度上升. 以使得在更小尺度生成的模式能被融合进更高尺度里，并填充一些细节.import time start = time.time() OCTAVE_SCALE = 1.30 img = tf.constant(np.array(original_img)) base_shape = tf.shape(img)[:-1] float_base_shape = tf.cast(base_shape, tf.float32) for n in range(-2, 3): new_shape = tf.cast(float_base_shape*(OCTAVE_SCALE**n), tf.int32) img = tf.image.resize(img, new_shape).numpy() img = run_deep_dream_simple(img=img, steps=50, step_size=0.01) # img = tf.image.resize(img, base_shape) img = tf.image.convert_image_dtype(img/255.0, dtype=tf.uint8) show(img) end = time.time() end-start8. 再升级 - 图片分片需要考虑的一点是，当图片尺寸增加时，梯度计算会需要更多的时间和内存消耗. 上面的实现可能就行不通了.对此，可以将图片进行分片，并分别计算每一个分片的梯度.为了避免图像分片连接处出现不协调，在分片计算前，对图像随机平移.def random_roll(img, maxroll): # Randomly shift the image to avoid tiled boundaries. shift = tf.random.uniform(shape=[2], minval=-maxroll, maxval=maxroll, dtype=tf.int32) img_rolled = tf.roll(img, shift=shift, axis=[0,1]) return shift, img_rolled # shift, img_rolled = random_roll(np.array(original_img), 512) show(img_rolled)图片分片后的 deepdream 函数：class TiledGradients(tf.Module): def __init__(self, model): self.model = model @tf.function( input_signature=( tf.TensorSpec(shape=[None,None,3], dtype=tf.float32), tf.TensorSpec(shape=[], dtype=tf.int32),) ) def __call__(self, img, tile_size=512): shift, img_rolled = random_roll(img, tile_size) # Initialize the image gradients to zero. gradients = tf.zeros_like(img_rolled) # Skip the last tile, unless there's only one tile. xs = tf.range(0, img_rolled.shape[0], tile_size)[:-1] if not tf.cast(len(xs), bool): xs = tf.constant([0]) ys = tf.range(0, img_rolled.shape[1], tile_size)[:-1] if not tf.cast(len(ys), bool): ys = tf.constant([0]) for x in xs: for y in ys: # Calculate the gradients for this tile. with tf.GradientTape() as tape: # This needs gradients relative to `img_rolled`. # `GradientTape` only watches `tf.Variable`s by default. tape.watch(img_rolled) # Extract a tile out of the image. img_tile = img_rolled[x:x+tile_size, y:y+tile_size] loss = calc_loss(img_tile, self.model) # Update the image gradients for this tile. gradients = gradients + tape.gradient(loss, img_rolled) # Undo the random shift applied to the image and its gradients. gradients = tf.roll(gradients, shift=-shift, axis=[0,1]) # Normalize the gradients. gradients /= tf.math.reduce_std(gradients) + 1e-8 return gradients # get_tiled_gradients = TiledGradients(dream_model)主循环：def run_deep_dream_with_octaves(img, steps_per_octave=100, step_size=0.01, octaves=range(-2,3), octave_scale=1.3): base_shape = tf.shape(img) img = tf.keras.preprocessing.image.img_to_array(img) img = tf.keras.applications.inception_v3.preprocess_input(img) initial_shape = img.shape[:-1] img = tf.image.resize(img, initial_shape) for octave in octaves: # Scale the image based on the octave new_size = tf.cast(tf.convert_to_tensor(base_shape[:-1]), tf.float32)*(octave_scale**octave) img = tf.image.resize(img, tf.cast(new_size, tf.int32)) for step in range(steps_per_octave): gradients = get_tiled_gradients(img) img = img + gradients*step_size img = tf.clip_by_value(img, -1, 1) if step % 10 == 0: display.clear_output(wait=True) show(deprocess(img)) print ("Octave {}, Step {}".format(octave, step)) result = deprocess(img) return result # img = run_deep_dream_with_octaves(img=original_img, step_size=0.01) display.clear_output(wait=True) img = tf.image.resize(img, base_shape) img = tf.image.convert_image_dtype(img/255.0, dtype=tf.uint8) show(img)9. 材料[1] - TensorFlow Lucid