<commit_before>"""Implementation of Hessian-free optimization.

Author: <NAME> (<EMAIL>)

Based on
<NAME>. (2010). Deep learning via Hessian-free optimization. In Proceedings
of the 27th International Conference on Machine Learning.

and Matlab code from <NAME> available at
http://www.cs.toronto.edu/~jmartens/research.html
"""

import pickle

from scipy.special import expit

import numpy as np

try:
    import pycuda.autoinit
    import pycuda.driver as drv
    from pycuda.compiler import SourceModule
    from pycuda import gpuarray
except ImportError:
    print "PyCuda not installed, cannot use GPU acceleration"


class HessianBackprop(object):
    def __init__(self, layers=[1, 1, 1], use_GPU=False, debug=False):
        self.use_GPU = use_GPU
        self.debug = debug
        self.n_layers = len(layers)
        self.layers = layers
        self.n_params = [0 for _ in range(self.n_layers - 1)]
        for i in range(self.n_layers - 1):
            self.n_params[i] += (self.layers[i] + 1) * self.layers[i + 1]

        self.init_weights()

        if use_GPU:
            self.init_GPU()
        else:
            def outer_sum(a, b, out):
                out[:] = np.ravel(np.einsum("ij,ik", a, b))
            self.outer_sum = outer_sum

    def init_GPU(self):
        dev = pycuda.autoinit.device
        print "GPU found, using %s %s" % (dev.name(), dev.compute_capability())
        self.num_threads = (np.int32(1024)
                            if dev.compute_capability()[0] >= 2 else
                            np.int32(512))

        # this one operation in the gradient/Gv calculations is where
        # most of the computational work can be this algorithm, so
        # parallelizing it on the GPU can be pretty helpful.
        mod = SourceModule("""
        __global__ void outer_sum(float *a, float *b, float *out,
                                  int batch_size)
        {
            int a_i = blockIdx.x*blockDim.x + threadIdx.x;
            int b_i = threadIdx.y;
            const int a_len = blockDim.x * gridDim.x;
            const int b_len = blockDim.y;
            const int out_addr = a_i*b_len + b_i;

            out[out_addr] = 0;
            for(int j=0; j < batch_size; j++) {
                out[out_addr] += a[a_i] * b[b_i];
                a_i += a_len;
                b_i += b_len;
            }
        }
        """)

        def outer_sum(a, b, out):
            a_len = np.int32(a.shape[1])
            b_len = np.int32(b.shape[1])
            batchsize = np.int32(a.shape[0])  # assume == b.shape[0]

            assert a.dtype == b.dtype == out.dtype == np.float32
            assert b_len < self.num_threads
            # note: could make this work with longer b's if needed,
            # would just need to tile horizontally in the same way as
            # is done vertically

            rows_per_block = np.minimum(self.num_threads / b_len,
                                        a_len)
            while a_len % rows_per_block != 0:
                rows_per_block -= 1

            gpu_outer = mod.get_function("outer_sum")
            a_gpu = (a if isinstance(a, gpuarray.GPUArray) else drv.In(a))
            b_gpu = (b if isinstance(b, gpuarray.GPUArray) else drv.In(b))
            out_gpu = (out if isinstance(out, gpuarray.GPUArray) else
                       drv.Out(out))
            gpu_outer(a_gpu, b_gpu, out_gpu, batchsize,
                      grid=(a_len / rows_per_block, 1),
                      block=(rows_per_block, b_len, 1))

            if self.debug:
                if isinstance(a, gpuarray.GPUArray):
                    tmp_a = np.zeros(a.shape, dtype=np.float32)
                    a.get(tmp_a)
                else:
                    tmp_a = a
                if isinstance(b, gpuarray.GPUArray):
                    tmp_b = np.zeros(b.shape, dtype=np.float32)
                    b.get(tmp_b)
                else:
                    tmp_b = b
                truth = np.ravel(np.einsum("ij,ik", tmp_a, tmp_b))
                try:
                    assert np.allclose(out, truth, atol=1e-6)
                except AssertionError:
                    print out
                    print truth
                    print out - truth
                    raise

        self.outer_sum = outer_sum

    def init_weights(self):
        """Weight initialization"""

        # sparse initialization (from martens)
        num_conn = 15
        W = [np.zeros((self.layers[i] + 1, self.layers[i + 1]),
                      dtype=np.float32)
             for i in range(self.n_layers - 1)]
        for l in range(self.n_layers - 1):
            for j in range(self.layers[l + 1]):
                # pick num_conn random pre neurons
                indices = np.ceil(np.random.rand(num_conn) * self.layers[l])
                indices = indices.astype(int)

                # connect to post
                W[l][indices, j] = np.random.randn(num_conn)
        self.W = np.concatenate([w.flatten() for w in W])

        # random initialization
#         self.W = np.zeros(np.sum(self.n_params), dtype=np.float32)
#         for i in range(self.n_layers - 1):
#             offset = np.sum(self.n_params[:i])
#             self.W[offset:offset + self.n_params[i]] = (
#                 np.random.uniform(-1 / np.sqrt(self.layers[i]),
#                                   1 / np.sqrt(self.layers[i]),
#                                   self.n_params[i]))

    def get_layer(self, params, layer, separate=True):
        """Get weight matrix for a layer from the overall parameter vector."""

        offset = np.sum(self.n_params[:layer])
        l_params = params[offset:offset + self.n_params[layer]]
        if separate:
            W = l_params[:-self.layers[layer + 1]]
            b = l_params[-self.layers[layer + 1]:]

            return (W.reshape((self.layers[layer],
                               self.layers[layer + 1])),
                    b)
        else:
            return l_params.reshape((self.layers[layer] + 1,
                                     self.layers[layer + 1]))

    def forward(self, input, params):<commit_msg>Modify weight precision when debugging<commit_after>"""Implementation of Hessian-free optimization.

Author: <NAME> (<EMAIL>)

Based on
<NAME>. (2010). Deep learning via Hessian-free optimization. In Proceedings
of the 27th International Conference on Machine Learning.

and Matlab code from <NAME> available at
http://www.cs.toronto.edu/~jmartens/research.html
"""

import pickle

from scipy.special import expit

import numpy as np

try:
    import pycuda.autoinit
    import pycuda.driver as drv
    from pycuda.compiler import SourceModule
    from pycuda import gpuarray
except ImportError:
    print "PyCuda not installed, cannot use GPU acceleration"


class HessianBackprop(object):
    def __init__(self, layers=[1, 1, 1], use_GPU=False, debug=False):
        self.use_GPU = use_GPU
        self.debug = debug
        self.n_layers = len(layers)
        self.layers = layers
        self.n_params = [0 for _ in range(self.n_layers - 1)]
        for i in range(self.n_layers - 1):
            self.n_params[i] += (self.layers[i] + 1) * self.layers[i + 1]

        self.init_weights()

        if use_GPU:
            self.init_GPU()
        else:
            def outer_sum(a, b, out):
                out[:] = np.ravel(np.einsum("ij,ik", a, b))
            self.outer_sum = outer_sum

    def init_GPU(self):
        dev = pycuda.autoinit.device
        print "GPU found, using %s %s" % (dev.name(), dev.compute_capability())
        self.num_threads = (np.int32(1024)
                            if dev.compute_capability()[0] >= 2 else
                            np.int32(512))

        # this one operation in the gradient/Gv calculations is where
        # most of the computational work can be this algorithm, so
        # parallelizing it on the GPU can be pretty helpful.
        mod = SourceModule("""
        __global__ void outer_sum(float *a, float *b, float *out,
                                  int batch_size)
        {
            int a_i = blockIdx.x*blockDim.x + threadIdx.x;
            int b_i = threadIdx.y;
            const int a_len = blockDim.x * gridDim.x;
            const int b_len = blockDim.y;
            const int out_addr = a_i*b_len + b_i;

            out[out_addr] = 0;
            for(int j=0; j < batch_size; j++) {
                out[out_addr] += a[a_i] * b[b_i];
                a_i += a_len;
                b_i += b_len;
            }
        }
        """)

        def outer_sum(a, b, out):
            a_len = np.int32(a.shape[1])
            b_len = np.int32(b.shape[1])
            batchsize = np.int32(a.shape[0])  # assume == b.shape[0]

            assert a.dtype == b.dtype == out.dtype == np.float32
            assert b_len < self.num_threads
            # note: could make this work with longer b's if needed,
            # would just need to tile horizontally in the same way as
            # is done vertically

            rows_per_block = np.minimum(self.num_threads / b_len,
                                        a_len)
            while a_len % rows_per_block != 0:
                rows_per_block -= 1

            gpu_outer = mod.get_function("outer_sum")
            a_gpu = (a if isinstance(a, gpuarray.GPUArray) else drv.In(a))
            b_gpu = (b if isinstance(b, gpuarray.GPUArray) else drv.In(b))
            out_gpu = (out if isinstance(out, gpuarray.GPUArray) else
                       drv.Out(out))
            gpu_outer(a_gpu, b_gpu, out_gpu, batchsize,
                      grid=(a_len / rows_per_block, 1),
                      block=(rows_per_block, b_len, 1))

            if self.debug:
                if isinstance(a, gpuarray.GPUArray):
                    tmp_a = np.zeros(a.shape, dtype=np.float32)
                    a.get(tmp_a)
                else:
                    tmp_a = a
                if isinstance(b, gpuarray.GPUArray):
                    tmp_b = np.zeros(b.shape, dtype=np.float32)
                    b.get(tmp_b)
                else:
                    tmp_b = b
                truth = np.ravel(np.einsum("ij,ik", tmp_a, tmp_b))
                try:
                    assert np.allclose(out, truth, atol=1e-6)
                except AssertionError:
                    print out
                    print truth
                    print out - truth
                    raise

        self.outer_sum = outer_sum

    def init_weights(self, dtype=np.float32):
        """Weight initialization"""
        if self.debug and dtype != np.float64:
            print "Changing weights to 64bit precision for debugging"
            dtype = np.float64

        # sparse initialization (from martens)
        num_conn = 15
        W = [np.zeros((self.layers[i] + 1, self.layers[i + 1]),
                      dtype=dtype)
             for i in range(self.n_layers - 1)]
        for l in range(self.n_layers - 1):
            for j in range(self.layers[l + 1]):
                # pick num_conn random pre neurons
                indices = np.ceil(np.random.rand(num_conn) * self.layers[l])
                indices = indices.astype(int)

                # connect to post
                W[l][indices, j] = np.random.randn(num_conn)
        self.W = np.concatenate([w.flatten() for w in W])

        # random initialization
#         self.W = np.zeros(np.sum(self.n_params), dtype=dtype)
#         for i in range(self.n_layers - 1):
#             offset = np.sum(self.n_params[:i])
#             self.W[offset:offset + self.n_params[i]] = (
#                 np.random.uniform(-1 / np.sqrt(self.layers[i]),
#                                   1 / np.sqrt(self.layers[i]),
#                                   self.n_params[i]))

    def get_layer(self, params, layer, separate=True):
        """Get weight matrix for a layer from the overall parameter vector."""

        offset = np.sum(self.n_params[:layer])
        l_params = params[offset:offset + self.n_params[layer]]
        if separate:
            W = l_params[:-self.layers[layer + 1]]
            b = l_params[-self.layers[layer + 1]:]

            return (W.reshape((self.layers[layer],
                               self.layers[layer + 1])),
                    b)
        else:
            return l_params.reshape((self.layers[layer] + 1,
                                     self.layers[layer + 1]))

    def forward(self, input, params):
{'fraction_non_alphanumeric': 0.10773059324074817, 'fraction_numerical': 0.011559765593642128, 'mean_word_length': 1.5766287487073423, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3664532', 'n_tokens_mistral': 3932, 'n_tokens_neox': 3687, 'n_words': 935}