нет ускорения при использовании numba jit - PullRequest
Новая
       26

нет ускорения при использовании numba jit

0 голосов
/ 06 мая 2020

Я пытаюсь запустить алгоритм NEAT, используя эту python реализацию. Это исходный файл из библиотеки, который имеет отношение к моему вопросу: 

from neat.graphs import feed_forward_layers


class FeedForwardNetwork(object):
    def __init__(self, inputs, outputs, node_evals):
        self.input_nodes = inputs
        self.output_nodes = outputs
        self.node_evals = node_evals
        self.values = dict((key, 0.0) for key in inputs + outputs)

    def activate(self, inputs):
        if len(self.input_nodes) != len(inputs):
            raise RuntimeError("Expected {0:n} inputs, got {1:n}".format(len(self.input_nodes), len(inputs)))

        for k, v in zip(self.input_nodes, inputs):
            self.values[k] = v

        for node, act_func, agg_func, bias, response, links in self.node_evals:
            node_inputs = []
            for i, w in links:
                node_inputs.append(self.values[i] * w)
            s = agg_func(node_inputs)
            self.values[node] = act_func(bias + response * s)

        return [self.values[i] for i in self.output_nodes]

    @staticmethod
    def create(genome, config):
        """ Receives a genome and returns its phenotype (a FeedForwardNetwork). """

        # Gather expressed connections.
        connections = [cg.key for cg in genome.connections.values() if cg.enabled]

        layers = feed_forward_layers(config.genome_config.input_keys, config.genome_config.output_keys, connections)
        node_evals = []
        for layer in layers:
            for node in layer:
                inputs = []
                node_expr = [] # currently unused
                for conn_key in connections:
                    inode, onode = conn_key
                    if onode == node:
                        cg = genome.connections[conn_key]
                        inputs.append((inode, cg.weight))
                        node_expr.append("v[{}] * {:.7e}".format(inode, cg.weight))

                ng = genome.nodes[node]
                aggregation_function = config.genome_config.aggregation_function_defs.get(ng.aggregation)
                activation_function = config.genome_config.activation_defs.get(ng.activation)
                node_evals.append((node, activation_function, aggregation_function, ng.bias, ng.response, inputs))

        return FeedForwardNetwork(config.genome_config.input_keys, config.genome_config.output_keys, node_evals)

Поскольку я оцениваю производительность своих нейронных сетей на большом наборе данных, я хотел ускорить метод activate с помощью numba джит. Чтобы не вернуться в объектный режим numbas, мне пришлось обновить реализацию метода activate (и, следовательно, также поля класса FeedForwardNetwork), используя только типы данных, поддерживаемые numba. Вот что я придумал (create то же, что и раньше): 

from neat.graphs import feed_forward_layers
from neat.six_util import itervalues

import numba
from numba import jit, njit
from numba.typed import List, Dict
import numpy as np
import math


@jit(nopython=True)
def activate(input_nodes, output_nodes, node_evals_node, node_evals_bias, node_evals_resp, node_evals_ins_nodes, node_evals_ins_conns, values, inputs):
    for i in range(input_nodes.size):
        values[input_nodes[i]] = inputs[i]

    for node in range(len(node_evals_node)):
        s = 0
        for pred in range(len(node_evals_ins_nodes[node])):
            s += values[node_evals_ins_nodes[node][pred]] * node_evals_ins_conns[node][pred]

        values[node_evals_node[node]] = math.tanh(node_evals_bias[node] + node_evals_resp[node] * s)

    return [values[output_nodes[i]] for i in range(output_nodes.size)]


class FeedForwardNetwork(object):
    def __init__(self, inputs, outputs, node_evals):
        self.input_nodes = np.array(inputs)
        self.output_nodes = np.array(outputs)

        # NODE_EVALS decomposition
        self.node_evals_node = np.reshape(np.array(node_evals)[:, 0:1], (len(node_evals),)).astype(np.int64)
        self.node_evals_bias = np.reshape(np.array(node_evals)[:, 3:4], (len(node_evals),)).astype(np.float64)
        self.node_evals_resp = np.reshape(np.array(node_evals)[:, 4:5], (len(node_evals),)).astype(np.float64)

        temp = np.array(node_evals)[:, 5:6]

        self.node_evals_ins_nodes = List()
        self.node_evals_ins_conns = List()

        for node in range(temp.size):
            l = List()
            m = List()

            for predecessor in range(len(temp[node])):
                l.append(temp[0][node][predecessor][0])
                m.append(temp[0][node][predecessor][1])

            self.node_evals_ins_nodes.append(l)
            self.node_evals_ins_conns.append(m)

        self.values = Dict()
        # Set types of dict
        self.values[0] = float(1)
        self.values.pop(0)

Это код, который я вызываю для методов create и activate в: 

def eval_single_genome(genome, config, thread_id, result):
    net = neat.nn.FeedForwardNetwork.create(genome, config)

    error_sum = 0

    for i, row in PRICES.iterrows():
        prediction = feed_forward.activate(net.input_nodes, net.output_nodes, net.node_evals_node, net.node_evals_bias, net.node_evals_resp, net.node_evals_ins_nodes, net.node_evals_ins_conns, net.values, np.array([0]))
        error_sum += (prediction - PRICES.iloc[i]['open']) ** 2

    result[thread_id] = error_sum

Код компилируется и запускается без ошибок или предупреждений, что (насколько я понял) указывает на то, что numba должна иметь возможность оптимизировать мою реализацию. Но добавление / удаление декоратора @jit(nopython=True) никак не меняет время выполнения.

Я что-то упустил? Или в моем случае numba ничего не может улучшить?

...