自定义报错返回_Keras编写自定义层--以GroupNormalization为例

1. Group Normalization 介绍

Batch Normalization(BN)称为批量归一化，可加速网络收敛利于网络训练。但BN的误差会随着批量batch的减小而迅速增大。FAIR 研究工程师吴育昕和研究科学家何恺明合作的一篇论文提出了一种新的与批量无关的Normalization 方法-[[1803.08494] Group Normalization]。GN 的主要工作是将通道分成组，并在每组内计算归一化的均值和方差。GN 的计算与批量大小无关，并且其准确度在各种批量大小下都很稳定。具体如下图(摘自论文)：

BN 时的小批量会导致批量数据的统计两估算不准确，会显著增加模型误差。而无批量无关的GN方法得到的误差则相对稳定。

2. Keras自定义层方法

关于Keras中如何自定义层，可参考官方中文文档[编写你自己的层 - Keras 中文文档 ]，[Keras简单自定义层例子]。自定义层中主要包括4种方法：

__init__(**kwargs):初始化方法，关键字参数保留，否则自定义层加载会报错。
build(input_shape)：用于定义权重的方法
call(x): 自定义层具体功能的实现方法
get_config : 返回一个字典，获取当前层的参数信息。自定义层保存和加载时需要定义
compute_output_shape(input_shape)：用于Keras可以自动推断shape

自定义层的保存和加载需要注意以下3点：

__init__(self, arg, **kwargs)初始化方法中关键字参数保留，否则自定义层加载会报错。

缺少**kwargs，TypeError: __init__() got an unexpected keyword argument 'name'

get_config(self)方法需要重写，否则网络结构无法保存。父类的config也需一并保存，将父类及继承类的config组装为字典形式，继承类config依据__init__方法传入的参数而定，具体如下：

缺少get_config方法，NotImplementedError: Layers with arguments in `__init__` must override `get_config`.

def get_config(self):base_config = super(LayerName, self).get_config() #父类config字典base_config['arg'] = self.arg #继承类config字典，__init__传入参数argreturn base_config #返回组装后的字典

load_model()需为custom_objects参数赋值

缺少custom_objects，ValueError: Unknown layer: LayerName

_custom_objects = {"LayerName":LayerName} #定义custom_objects
model = keras.models.load_model(model_path, custom_objects=_custom_objects) #加载模型

Keras官网提供了两种Normalization的源码，分别是：

批量归一化 Keras-BatchNormalization
实例归一化 Keras-InstanceNormalization

两者的不同在于IN的统计量估算是批量无关的基于单张图片单个通道，不需要用滑动平均项来记录全局的统计量，体现在源码的差异为：

# BN code
class BatchNormalization(Layer):def __init__(self,**kwargs)：super(BatchNormalization, self).__init__(**kwargs)...def build(self, input_shape):...self.gamma = self.add_weight(...)self.beta = self.add_weight(...)self.moving_mean = self.add_weight(trainable=False)self.moving_variance = self.add_weight(trainable=False)...def call(sekf, x):...self.add_update([K.moving_average_update(self.moving_mean, mean,self.momentum),K.moving_average_update(self.moving_variance,variance,self.momentum)],inputs)...return K.in_train_phase(...)

# IN code
class BatchNormalization(Layer):def __init__(self,**kwargs)：super(InstanceNormalization, self).__init__(**kwargs)...def build(self, input_shape):...self.gamma = self.add_weight(...)self.beta = self.add_weight(...)...

解释：

所有自定义层都需要继承基础层Layer，并添加super().__init__(**kwargs)
**kwargs代表以字典方式继承父类
self.add_weight()是继承层Layer的方法，用于为变量添加权重，其中有参数trainable代表该参数的权重是否为可训练权重; 若trainable==True时，会执行self._trainable_weights.append(weight).
BN中需要添加moving_mean/variance滑动平均项的权重，且需要设置trainable==False，即为非训练参数。
self.add_update()用于更新滑动平均项
K.in_train_phase()针对训练状态选择不同的mean/variance计算BN

3. 定义Group Normalization层

源代码位置Bingohong/GroupNormalization-tensorflow-keras，里面包含了2个GN文件，分别是tensorflow和keras的实现版本，其中都包含了moving_average操作。

其实关于GN操作，是否需要apply moving_average是值得商榷的，论文中貌似没有明确提及，其他实现版本中的实现都是无moving_average操作。但通过对比IN、BN和GN特点及后期的实验对比，觉得GN应该是不需要moving_average操作的。因此这部分内容包括：

主要介绍有moving_average操作的GN层的定义过程，而无moving_average操作时，只需要将对应的代码去掉。
使用BN/GN_with_moving_average/GN_without_moving_average3种Normalization方法，对比U-net的实验结果。

keras GN层

完整代码在这里，以下仅解释部分关键代码。

# GN_with_moving_average code
class GroupNormalization(Layer):def __init__(self,**kwargs)：super(GroupNormalization, self).__init__(**kwargs)...def build(self, input_shape):...shape = (self.groups,)broadcast_shape = [-1, self.groups, 1, 1, 1]# 添加滑动平均项参数，并设置为非训练参数# 后续的K.reshape和K.variable操作，是为了在call()方法内进行add_update()时保证# self.moving_mean/variance的维度与inputs一致，且为variable变量self.moving_mean = self.add_weight(shape=shape,trainable=False)self.moving_mean = K.reshape(self.moving_mean,broadcast_shape)self.moving_mean = K.variable(value=self.moving_mean)self.moving_variance = self.add_weight(shape=shape,trainable=False)self.moving_variance = K.reshape(self.moving_variance,broadcast_shape)self.moving_variance = K.variable(value=self.moving_variance)...def call(sekf, inputs):G = self.groups# transpose:[ba,h,w,c] -> [bs,c,h,w]if self.axis in {-1,3}:inputs = K.permute_dimensions(inputs,(0,3,1,2))# GN操作需要根据groups对通道分组input_shape = K.int_shape(inputs)N, C, H, W = input_shapeinputs = K.reshape(inputs,(-1, G, C // G, H, W))#inputs.assign_sub()# 计算分组通道的均值和方差gn_mean = K.mean(inputs,axis=[2,3,4],keepdims=True)gn_variance = K.var(inputs,axis=[2,3,4],keepdims=True)# 当模型用于测试阶段时，使用moving_mean/variance记录的均值/方差def gn_inference():# when in test phase, just return moving_mean & moving_varmean, variance = self.moving_mean, self.moving_varianceoutputs = (inputs - mean) / (K.sqrt(variance + self.epsilon))outputs = K.reshape(outputs,[-1, C, H, W]) * self.gamma + self.beta# transpose: [bs,c,h,w] -> [ba,h,w,c]if self.axis in {-1,3}:outputs = K.permute_dimensions(outputs,(0,2,3,1))return outputsif training in {0,False}:return gn_inference()# 当模型用于训练阶段时，使用分组通道实时计算均值/方差outputs = (inputs - gn_mean) / (K.sqrt(gn_variance + self.epsilon))outputs = K.reshape(outputs,[-1, C, H, W]) * self.gamma + self.beta # transpose: [bs,c,h,w] -> [ba,h,w,c]if self.axis in {-1,3}:outputs = K.permute_dimensions(outputs,(0,2,3,1))# 手动更新self.moving_mean/varianceself.add_update([K.moving_average_update(self.moving_mean, mean,self.momentum),K.moving_average_update(self.moving_variance,variance,self.momentum)],inputs)# 根据模型状态不同选择不同的GN计算方法，train时选择outputs，test时选择gn_inferencereturn K.in_train_phase(outputs, gn_inference,training=training)

实验对比结果

实验日志位于compare_log，包含3个文件：

train_bn.log -> unet+bn日志
train_gn_ema.log -> unet+gn(有moving_average操作)
train_gn_noema.log -> unet+gn(无moving_average操作)

结果说明：

gn without moving average 得到的val_loss会更低，可达到 0.2左右
gn with moving average 有时会一直存在很高的val_loss, 所以我觉得可能GN并不需要 apply moving average
bn 得到的val_ loss约为 0.26, 高于gn without moving average.

欢迎大家批评指正～谢谢谢谢～～～