前言
前面写了一篇利用TF在iOS实现类Prisma软件的文章后，收到很多网上朋友交流实现思想与求Demo回复，可见大家对于这个功能实现有浓厚的兴趣。
上一篇文章并未深入详解Google的实现的原理，仅仅是简单将参数和计算图在iOS设备上跑起来，且由于TF本身编译搭建工程就很复杂，所以未托管源码供下载。
本次主要是通过剖析Google论文中实现逻辑后，利用iOS新的Metal框架以及设备的GPU部分加速运算后实现了图片的类Prisma渲染（当然网络参数还是Google训练出来的）。期间跨过了很多坑，这里分享出来，希望对大家研究有所帮助😅。
2017草莓音乐节(阴阳师舞台)
结构
看过《A LEARNED REPRESENTATION FOR ARTISTIC STYLE》这篇论文的应该都还是比较熟悉Google提出的整体网络结构，如下图：
网络结构
后面虚线部分是VGG网络，这部分训练方法和2015年《A Neural Algorithm of Artistic Style》这篇论文中一样，并非google此次优化的重点。
重点在前面Style transfer network这部分，我们的到的训练参数也都是这一层的参数。这是一个前向生成图片的网络，有了这么一个前向网络层，我们只需要训练好参数，生成图片就只需前向运算一遍就可以的到，这样相对于直接利用VGG网络来回学习生成图片节省了很大一部分时间。并且由于合成图片时间短，所以在移动设备上也是可以本地运行。
下面是Style transfer network的网络结构：
network
网络一共3个卷积+5个Residual Block+2个upsampling+1个卷积，其实Residual Block就是两次卷积，然后将输入与输出相加，upsampling是先用Nearest-Neighbor放大图片，然后卷积。所以一共有16层卷积操作，并且每次卷积操作后先Batch-normalization，再接激活函数（这里开始因为苹果MPSCNN库直接就可以卷积后带激活函数，所以在实现的时候我把BN放到了激活之后，生成图片就一直错误😭）。
Padding Mode 论文这里写的是Reflect，苹果Metal不支持这个padding方式，我自己写了一个😢，但最后发现其实用Zero Padding反而才是对的。不知道是否是苹果的卷积实现有不同，还是说这里Padding模式仅仅是针对训练的时候？这里后面有时间会再研究一下。
以上就是整个实现的核心网络结构，理论上我们有了参数，知道了网络实现，不用TF计算图，自己实现也是可以的。这样可以免去繁琐的TF集成，编译，并且自己的网络调试、控制内存等等都要方便很多。
但是，并非如此简单，苹果Metal框架很多深度神经网络的kernel都还没有，仅仅对卷积操作有部分封装。下面就分享实现过程中几个比较重要算法的实现。
Batch-Normalization
BN(Batch-Normalization)其实是这个网络的核心部分，每种不同的style图片就是在这里进行区分的，当你选择不同的style的时候，每层卷积操作是相同的，但是BN不同，就是改变最后生成图片的样式。
开始我期望Metal有BN的实现，但找了一圈没发现，考虑过写kernel这样和卷积操作都能在GPU上运算，但是最后发现kernel编码从头学太复杂，于是就在CPU上实现一个，每次卷积完过后，图片Copy出来在CPU上运算BN，然后再接激活函数（还是期待苹果后面能提供支持BN的kernel😊）。
实现代码如下：
- (void)batch_norm:(MPSImage *)image styles:(float *)styles shift:(float *)shift
{
NSUInteger w = image.texture.width;
NSUInteger h = image.texture.height;
NSUInteger featureNum = image.featureChannels;
float *gamma = calloc(featureNum, sizeof(float));
float *beta = calloc(featureNum, sizeof(float));
// float gamma[featureNum], beta[featureNum];
vDSP_mmul(styles, 1, shift, 1, beta, 1, 1, featureNum, styleNum);
vDSP_mmul(styles, 1, shift+featureNum*styleNum, 1, gamma, 1, 1, featureNum, styleNum);
// for (int i = 0; i < featureNum; i++) {
// printf("%f,%f ",gamma[i],beta[i]);
// }
// NSLog(@"%@",image);
//
NSUInteger numSlices = (featureNum + 3) / 4;
NSUInteger numComponents = featureNum < 3 ? featureNum : 4;
NSUInteger channels = featureNum < 3 ? featureNum : numSlices * 4;
float16_t *htemp = calloc(w*h*channels, sizeof(float16_t));
for (int i = 0; i < numSlices; i++) {
[image.texture getBytes:htemp+w*h*numComponents*i bytesPerRow:w*numComponents*2 bytesPerImage:0 fromRegion:MTLRegionMake3D(0, 0, 0, w, h, 1) mipmapLevel:0 slice:i];
}
float *temp = calloc(w*h*channels, sizeof(float));
[self halfTofloat:htemp floatp:temp width:w height:h channel:channels];
float mean, var;
for (int i = 0; i < featureNum; i++) {
int slice = i / 4;
int stride = i % 4;
vDSP_normalize(temp+slice*w*h*numComponents+stride, numComponents, temp+slice*w*h*numComponents+stride, numComponents, &mean, &var, w*h);
if (var == 0) {
vDSP_vfill(&var, temp+slice*w*h*numComponents+stride, numComponents, w*h);
}
vDSP_vsmul(temp+slice*w*h*numComponents+stride, numComponents, &gamma[i], temp+slice*w*h*numComponents+stride, numComponents, w*h);
vDSP_vsadd(temp+slice*w*h*numComponents+stride, numComponents, &beta[i], temp+slice*w*h*numComponents+stride, numComponents, w*h);
}
[self floatToHalf:temp halfp:htemp width:w height:h channel:channels];
for (int i = 0; i < numSlices; i++) {
[image.texture replaceRegion:MTLRegionMake3D(0, 0, 0, w, h, 1) mipmapLevel:0 slice:i withBytes:htemp+w*h*numComponents*i bytesPerRow:w*numComponents*2 bytesPerImage:0];
}
free(temp);
free(htemp);
free(gamma);
free(beta);
}
Nearest-Neighbor
这个填充算法，苹果也没有直接提供，BlitCommandEncoder里面有相关的方法，但是我感觉使用有点麻烦，本来是个很简单的填充算法，再加上前面BN都已经在CPU上实现了，这个也就调用2次，于是我也直接在CPU上去实现运算了。
原理很简单，就是放大图片像素点周围用这一个色值去填充，
Nearest-Neighbor
实现代码：
- (void)ResizeNearestNeighbor:(MPSImage *)source destinationImage:(MPSImage *)destinationImage
{
NSUInteger w = source.texture.width;
NSUInteger h = source.texture.height;
NSUInteger w2 = destinationImage.texture.width;
NSUInteger h2 = destinationImage.texture.height;
NSUInteger featureNum = source.featureChannels;
NSUInteger numSlices = (featureNum + 3) / 4;
NSUInteger numComponents = featureNum < 3 ? featureNum : 4;
NSUInteger channels = featureNum < 3 ? featureNum : numSlices * 4;
float16_t *htemp1 = calloc(w*h*channels, sizeof(float16_t));
float16_t *htemp2 = calloc(w2*h2*channels, sizeof(float16_t));
for (int i = 0; i < numSlices; i++) {
[source.texture getBytes:htemp1+w*h*numComponents*i bytesPerRow:w*numComponents*2 bytesPerImage:0 fromRegion:MTLRegionMake3D(0, 0, 0, w, h, 1) mipmapLevel:0 slice:i];
}
int x_ratio = (int)((w<<16)/w2) +1;
int y_ratio = (int)((h<<16)/h2) +1;
int x2, y2 ;
for (int k = 0; k < featureNum; k++) {
int slice = k / 4;
int stride = k % 4;
for (int i=0;i<h2;i++) {
for (int j=0;j<w2;j++) {
x2 = ((j*x_ratio)>>16) ;
y2 = ((i*y_ratio)>>16) ;
htemp2[slice*w2*h2*numComponents+(i*w2+j)*numComponents+stride] = htemp1[slice*w*h*numComponents+((y2*w)+x2)*numComponents+stride];
}
}
}
for (int i = 0; i < numSlices; i++) {
[destinationImage.texture replaceRegion:MTLRegionMake3D(0, 0, 0, w2, h2, 1) mipmapLevel:0 slice:i withBytes:htemp2+w2*h2*numComponents*i bytesPerRow:w2*numComponents*2 bytesPerImage:0];
}
free(htemp1);
free(htemp2);
}
整个网络实现
最后，整个网络的实现，参照论文里的结构和链接顺序，其中所有的卷积都是继承的MPSCNNConvolution对象，代码有点长如下：
- (MPSImage *)forward:(CGImageRef)srcImage width:(int)width height:(int)height styles:(float *)styles
{
id<MTLCommandBuffer> commandbuffer = [commandQueue commandBuffer];
int w = width;
int h = height;
MTKTextureLoader *loader = [[MTKTextureLoader alloc] initWithDevice:mtDevice];
id<MTLTexture> srcTexture = [loader newTextureWithCGImage:srcImage options:nil error:nil];
MPSImage *cc1Image = [[MPSImage alloc] initWithTexture:srcTexture featureChannels:3];
// MPSImage *tImage = [[MPSImage alloc] initWithTexture:srcTexture featureChannels:3];
// MPSImageDescriptor *cc1Des = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 width:w height:h featureChannels:3];
// MPSImage *cc1Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:cc1Des];
// [cc1Image.texture replaceRegion:MTLRegionMake3D(0, 0, 0, w, h, 1) mipmapLevel:0 withBytes:srcImage bytesPerRow:w*4*2];
// contract
MPSImageDescriptor *cc2Des = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 width:w height:h featureChannels:32];
MPSImage *cc2Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:cc2Des];
[contractConv1 encodeToCommandBuffer:commandbuffer sourceImage:cc1Image destinationImage:cc2Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:cc2Image styles:styles shift:cc1Shift];
commandbuffer = [commandQueue commandBuffer];
[relu encodeToCommandBuffer:commandbuffer sourceImage:cc2Image destinationImage:cc2Image];
w /= 2;
h /= 2;
MPSImageDescriptor *cc3Des = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 width:w height:h featureChannels:64];
MPSImage *cc3Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:cc3Des];
[contractConv2 encodeToCommandBuffer:commandbuffer sourceImage:cc2Image destinationImage:cc3Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:cc3Image styles:styles shift:cc2Shift];
commandbuffer = [commandQueue commandBuffer];
[relu encodeToCommandBuffer:commandbuffer sourceImage:cc3Image destinationImage:cc3Image];
w /= 2;
h /= 2;
MPSImageDescriptor *rcDes = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 width:w height:h featureChannels:128];
MPSImage *rc11Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[contractConv3 encodeToCommandBuffer:commandbuffer sourceImage:cc3Image destinationImage:rc11Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:rc11Image styles:styles shift:cc3Shift];
// residual
commandbuffer = [commandQueue commandBuffer];
[relu encodeToCommandBuffer:commandbuffer sourceImage:rc11Image destinationImage:rc11Image];
MPSImage *rc12Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[residual1Conv1 encodeToCommandBuffer:commandbuffer sourceImage:rc11Image destinationImage:rc12Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:rc12Image styles:styles shift:rc11Shift];
commandbuffer = [commandQueue commandBuffer];
[relu encodeToCommandBuffer:commandbuffer sourceImage:rc12Image destinationImage:rc12Image];
MPSImage *rc21Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[residual1Conv2 encodeToCommandBuffer:commandbuffer sourceImage:rc12Image destinationImage:rc21Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:rc21Image styles:styles shift:rc12Shift];
[self addImage:rc11Image B:rc21Image C:rc21Image];
commandbuffer = [commandQueue commandBuffer];
MPSImage *rc22Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[residual2Conv1 encodeToCommandBuffer:commandbuffer sourceImage:rc21Image destinationImage:rc22Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:rc22Image styles:styles shift:rc21Shift];
commandbuffer = [commandQueue commandBuffer];
[relu encodeToCommandBuffer:commandbuffer sourceImage:rc22Image destinationImage:rc22Image];
MPSImage *rc31Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[residual2Conv2 encodeToCommandBuffer:commandbuffer sourceImage:rc22Image destinationImage:rc31Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:rc31Image styles:styles shift:rc22Shift];
[self addImage:rc21Image B:rc31Image C:rc31Image];
commandbuffer = [commandQueue commandBuffer];
MPSImage *rc32Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[residual3Conv1 encodeToCommandBuffer:commandbuffer sourceImage:rc31Image destinationImage:rc32Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:rc32Image styles:styles shift:rc31Shift];
commandbuffer = [commandQueue commandBuffer];
[relu encodeToCommandBuffer:commandbuffer sourceImage:rc32Image destinationImage:rc32Image];
MPSImage *rc41Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[residual3Conv2 encodeToCommandBuffer:commandbuffer sourceImage:rc32Image destinationImage:rc41Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:rc41Image styles:styles shift:rc32Shift];
[self addImage:rc31Image B:rc41Image C:rc41Image];
commandbuffer = [commandQueue commandBuffer];
MPSImage *rc42Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[residual4Conv1 encodeToCommandBuffer:commandbuffer sourceImage:rc41Image destinationImage:rc42Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:rc42Image styles:styles shift:rc41Shift];
commandbuffer = [commandQueue commandBuffer];
[relu encodeToCommandBuffer:commandbuffer sourceImage:rc42Image destinationImage:rc42Image];
MPSImage *rc51Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[residual4Conv2 encodeToCommandBuffer:commandbuffer sourceImage:rc42Image destinationImage:rc51Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:rc51Image styles:styles shift:rc42Shift];
[self addImage:rc41Image B:rc51Image C:rc51Image];
commandbuffer = [commandQueue commandBuffer];
MPSImage *rc52Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[residual5Conv1 encodeToCommandBuffer:commandbuffer sourceImage:rc51Image destinationImage:rc52Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:rc52Image styles:styles shift:rc51Shift];
commandbuffer = [commandQueue commandBuffer];
[relu encodeToCommandBuffer:commandbuffer sourceImage:rc52Image destinationImage:rc52Image];
MPSImage *temp = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:rcDes];
[residual5Conv2 encodeToCommandBuffer:commandbuffer sourceImage:rc52Image destinationImage:temp device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:temp styles:styles shift:rc52Shift];
[self addImage:rc51Image B:temp C:temp];
// unsampling
commandbuffer = [commandQueue commandBuffer];
w *= 2;
h *= 2;
MPSImageDescriptor *ec1Des = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 width:w height:h featureChannels:128];
MPSImage *ec1Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:ec1Des];
[self ResizeNearestNeighbor:temp destinationImage:ec1Image];
MPSImageDescriptor *temp2Des = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 width:w height:h featureChannels:64];
MPSImage *temp2 = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:temp2Des];
[expandConv1 encodeToCommandBuffer:commandbuffer sourceImage:ec1Image destinationImage:temp2 device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:temp2 styles:styles shift:ec1Shift];
commandbuffer = [commandQueue commandBuffer];
[relu encodeToCommandBuffer:commandbuffer sourceImage:temp2 destinationImage:temp2];
w *= 2;
h *= 2;
MPSImageDescriptor *ec2Des = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 width:w height:h featureChannels:64];
MPSImage *ec2Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:ec2Des];
[self ResizeNearestNeighbor:temp2 destinationImage:ec2Image];
MPSImageDescriptor *ec3Des = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 width:w height:h featureChannels:32];
MPSImage *ec3Image = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:ec3Des];
[expandConv2 encodeToCommandBuffer:commandbuffer sourceImage:ec2Image destinationImage:ec3Image device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:ec3Image styles:styles shift:ec2Shift];
commandbuffer = [commandQueue commandBuffer];
[relu encodeToCommandBuffer:commandbuffer sourceImage:ec3Image destinationImage:ec3Image];
MPSImageDescriptor *destDes = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 width:w height:h featureChannels:3];
MPSImage *destImage = [[MPSImage alloc] initWithDevice:mtDevice imageDescriptor:destDes];
[expandConv3 encodeToCommandBuffer:commandbuffer sourceImage:ec3Image destinationImage:destImage device:mtDevice];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
[self batch_norm:destImage styles:styles shift:ec3Shift];
commandbuffer = [commandQueue commandBuffer];
[sigmoid encodeToCommandBuffer:commandbuffer sourceImage:destImage destinationImage:destImage];
[commandbuffer commit];
[commandbuffer waitUntilCompleted];
return destImage;
}
结语
什么话都不想留下了☠️，放几张程序运行图吧😊。
运行时内存、GPU、CPU状态
网友评论