实现在rasa-core中给policy提速
之前在rasa对话系统踩坑记(六)中提到过EmbeddingPolicy
和KerasPolicy
运行起来特别慢,原因是作者没有充分利用CPU资源,没有使用GPU资源,然后我对其进行了优化,并且也给rasa-core提交了pull request。之所以拖到现在才写这块的总结,因为期间一直在让Ghostvv给我review代码(不得不说德国人的认真劲让我吃了不少苦头),然后每次都是晚上他给我提修改建议,白天我改过来,所以比较慢,而且开始为了跑通Travis CI花了好久。到目前为止还没有被merge到rasa-core主线上,只能耐心等待。
具体的实现真的很简单,只是加了个tensorflow session config
的配置。做这个优化的初衷是因为我们实际训练story真的好慢好慢,而且还只是用的KerasPolicy
,慢到开始怀疑是不是服务器太烂。然后同组的飞龙同学发现rasa-core的policy没有支持GPU的配置,待加上了之后让本来需要半天才能跑完的模型半小时就结束了。
一开始我是在自定义的policy上将继承的policy的对应方法进行了重写,比如:
config_proto = {
"device_count": cpu_count(),
"inter_op_parallelism_threads": 0,
"intra_op_parallelism_threads": 0,
"allow_growth": True
}
def get_config_proto(config_proto_params):
# 配置configProto
config = tf.ConfigProto(
device_count={
'CPU': config_proto_params.get('device_count')
},
inter_op_parallelism_threads=config_proto_params.get('inter_op_parallelism_threads', 0),
intra_op_parallelism_threads=config_proto_params.get('intra_op_parallelism_threads', 0),
gpu_options={
'allow_growth': config_proto_params.get('allow_growth', True)
}
)
return config
class CustomPolicy(KerasPolicy):
def model_architecture(self, input_shape, output_shape):
"""Build a Keras model and return a compiled model."""
...
def train(self, training_trackers, domain, **kwargs):
...
self.graph = tf.Graph()
with self.graph.as_default():
self.session = tf.Session(config=get_config_proto(config_proto))
...
@classmethod
def load(cls, path):
...
graph = tf.Graph()
with graph.as_default():
session = tf.Session(config=get_config_proto(config_proto))
...
...
这里是继承自KerasPolicy
,并且将其中的train
和load
方法重置了,在tf.Session
加上了config
的参数配置。tf.ConfigProto
这里主要设置了4个参数,分别的含义是:
- device_count:告诉tf Session使用CPU数量上限,如果你的CPU数量较多,可以适当加大这个值
- inter_op_parallelism_threads: 控制运算符op内部的并行,当运算符op为单一运算符,并且内部可以实现并行时,如矩阵乘法,reduce_sum之类的操作,可以通过设置intra_op_parallelism_threads 参数来并行, intra代表内部。
- intra_op_parallelism_threads:控制多个运算符op之间的并行计算,当有多个运算符op,并且他们之间比较独立,运算符和运算符之间没有直接的路径Path相连。Tensorflow会尝试并行地计算他们,使用由inter_op_parallelism_threads参数来控制数量的一个线程池。
- allow_growth:如果设置为True,将动态申请显存
但是这样做就很麻烦,每次都要将KerasPolicy
方法重置。所以后面想到的是直接给官网提意见,提pr。和Ghostvv沟通后,需要在KerasPolicy
和EmbeddingPolicy
都要将其实现。看似很简单,只要在两个policy
的train
和load
方法上加上tf.ConfigProto
就行了,事实也只需要如此。但作为开源项目可不单单做这些就可以了。需要考虑统一性、复用性和简洁性。最终是在Policy
这个父类上加了
_load_tf_config
的静态方法和tf_defaults
静态属性,并在KerasPolicy
和EmbeddingPolicy
的_load_params
方法上分别对其parameters进行了merge。代码如下:
def _load_params(self, **kwargs: Dict[Text, Any]) -> None:
self.config = copy.deepcopy(dict(self.defaults, **self.tf_defaults))
self.config.update(kwargs)
self._tf_config = self._load_tf_config(self.config)
self._load_nn_architecture_params(self.config)
self._load_embedding_params(self.config)
self._load_regularization_params(self.config)
self._load_attn_params(self.config)
self._load_visual_params(self.config)
还有在数据存储这块,最终也是将_load_tf_config
返回的值进行了pickle的存储。省去了每次都要调用_load_tf_config
这个方法。细节这块现在还只能看pull request 1494,等合并到主线,就可以使用了。
很期待成为rasa-core的contributor,之前已经很水的成为rasa-nlu的contributor了。原创文章,转载请说明出处