如今人工智能大行其道， 其中在图像、影音处理中方法颇多。 本文介绍一种游戏中图像风格转换的例子，训练采用Tensorflow-GAN的方式，运行时在Unity引擎使用compute shader实现了跟tensorflow中一样的前向传播的网络来生成转换后的风格。对应Tensorflow环境中的Generator。 对应到github地址点击这里, 实现的效果如下图, 转换风格的图像我们通过一张RT渲染到右下角：

训练集

我们在python-tensorflow中实现了一种反向传播网络，我们使用Auto-Encoder替代GAN中Generator，用以生成风格化的图像， 而在Discriminator来鉴别图像。

训练中采用的训练集是微软的coco的dataset, 下载转换风格图片集.

我们在训练Discriminator的时候， 评估损失函数，丢给Discriminator风格化的图片尽量大， coco训练集的图片因为是假的， 我们去使其输出的值尽量小， 通过gennerator的图片，也是假的，我们也使其值也越小。

在训练Genenrator评估其损失函数的时候， generator的图片丢给Discriminator，尽量和真实的风格图片尽量接近，因此我们使其输出的值也要越大。

这里定义content image的内容的损失，方法采用的如paper 均值池化提取特征后， 比较方差的差异。

定义style featur的损失， 方法采用把generator生成的图片和原始输入的style image分别带入encoder做绝对值的差。

# Image loss.
self.img_loss_photo = mse_criterion(
    transformer_block(self.output_photo), transformer_block(self.input_photo))
self.img_loss = self.img_loss_photo

# Features loss.
self.feature_loss_photo = abs_criterion(self.output_photo_features, self.input_photo_features)
self.feature_loss = self.feature_loss_photo

Tensorboard 里观察loss的变化：

运行时

在unity中， 我们使用compute shader实现了一套跟TensorFlow中的相同的前向传播网络。 这里只是实现了Encoder和Decoder， 并没有实现Discriminator, 因为只有训练的时候用到了Discriminator.

在网络层中 Batch-normal由于为了求整体的均值和方差， 需要遍历当前层每个深度的layer，需要使用归纳算法-reduce, 使计算效率时间复杂度由n变成logn, 所以设计网络的时候尽量减少了类似的操作，在CompVis的设计网络中，decoder使用了九个残差模块， 为了性能我们这里减少到了一个。 参数规模也由大概48M减少到12M， 却实现了类似的效果。

由于受限于compute shader的语法， 我们在定义thread group时候， thread的大小不能超过1024，thread-z不能超过64， group的组成是vec3的格式， 而且需要是32或者64的倍数， 因此我们在设计网络的时候，每一个的layer尽量去靠近这些特性， 致使GPU发挥出最大的性能。 （CS5 group thread个数最多是1024， 而CS4最多支持到512，Apple的平台最多支持到CS4，这里需要注意下）。

compute shader里网络的参数规模如下图所示：

        

简化过后的，对应的tensorboard里的版本：

        

在设计的CNN网络层中， 卷积核大都是3x3的， 所以这里我们定义了一个Matrix3x3.cs在csharp中， 这样给StructureBuffer传递数据的时候，一次性传递过去。

数据导出

tensorflow 训练数据集有自己的序列化方式，大概是protobuf, google也提供了一套api, 去获取里面的张量Tensor。

通过训练集之后导出的checkpoint文件大小超过一个G，如果把这么庞大的参数文件导入到unity中所开销的内存空间是无法想象的。

通过遍历checkpoint发现所有的tensor发现， 每一层layer， 网络中的每个参数都对应了两个Adam对象，所以我们写了一套工具，导出数据的时候过滤掉Adam对象，使其大小减小到之前的1/3.

在上面提到，discriamtor只在训练的时候用到，而且其参数规模远超generator, 这里我们在导出的时候也需要对其过滤掉。

通过上面的操作我们导出的参数规模大概是48M，由于可以去掉预算复杂且效率不高的残差网络模块，参数规模进一步缩小到十几兆。 这还是我们采用float存储的格式， 如果对精度要求不高， 采用half的数据格式 5M-6M之间， 我觉得这个大小都手机平台还是可以接受的。

数据调试

为了清晰的看出Tensorflow Session里的每一层数据， 我们同时在python和unity环境中定义了printf 和 printh 函数， 分别用来输出二三维和一二维数据到控制台。 也可以检测python环境和unity环境运算结果是否具有一致性。

为了看出神经网络中的参数整体的相关性， 这里通过工具导出Auto-Encoder每一层里的数据（.bytes 二进制文件）， 然后在unity导入数据。通过Unity-Editor工具，我们对数据HXW前两维的数据生成一张Texture上， 在生成数据之前，我们通过激励函数sigmod, 使全部数据normalized到取件[0,1]之间。 之后，在通过一个slider表示第三维度的数据，通过拖拽slider的进度， 就可以看到不同depth的数据整体相关情况了。

我们通过一张图片看到， ganerator前期encoder前期更多的学到的是形状信息， 后期decoder学到的更多的是风格相关的信息。

性能分析

1. 参数规模:
   

在导出参数二进制文件到unity中的时候， 我们在python 控制台输出了参数规模大小约88M， 在过滤掉只在训练集里用到的参数， 导出的文件大概是9.7M。 在项目的fast分支， 我们删除了一些不重要的layer， 精简网络之后，参数规模占用23M， 过滤掉只在训练集里用到的参数， 导出的文件大小只有600K多一些了， 相比较于Tensorflow中Checkpoint文件中多达一个多G的参数规模，我们在运行时参数规模精简了1000多倍。（tensorflow checkpoint文件采用了protobuf格式的数据存储，我们这里纯是二进制， 因为protobuf虽然也是二进制， 但是会记一些tag之类的版本兼容相关的信息， 另外tensorflow还会存discriminator的信息，这块实现比较复杂，但只会出现在训练集里）。

1. 帧率:
   

当我们使用完整的pix2pixHD去实现的时候， PC上大概能跑3FPS, MAC上大概还不到1FPS。 之后在fast分支， 我们去掉了一些不重要的网络层，PC上大概能跑到29FPS，在文章的效果图里可以看到对应的效果。

在手机上， 我们对模糊效果的处理，往往为了性能会控制采样点的数量，也就是卷积核的大小。当采样点超过9个的话， 性能相对来说就会有一定的下降。 如果按此来比对基于深度神经网络方式实现的风格转换的做法时，那么这里采样的点岂止9个，即使fast上的实现方式，每个像素的采样规模也是模糊处理的几十倍， 所以帧率低也是有一定道理的。 在fast分支处理的时候， 我们尽量删除了那些depth标记深的layer, 因为cnn的时候要对每一层累加， 归纳reduce算法对于GPU来说，因为要遍历全局，时间复杂度最快也是logn。