rasa对话系统踩坑记（九）

实现在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了。原创文章，转载请说明出处