<jupyter_start><jupyter_text>Why is this interesting?Depending on how we caluclate gradients, some things are learnable while others are not. Linear networks (trained with SGD) have been known not to ...Given a better teacher, it is possible to ... The difference between model-based and model-free optimisation. The model-based optimisation (SGD -- where we use a linear model, ie the gradient) is very close to being accurate, but is not quite true. VS. Model-free optimisation doesnt care, it just looks at outputs and ...<jupyter_code>import numpy as np
import numpy.random as rnd

import torch as tr
import sklearn.datasets
dtype = tr.uint8
n_pop = 128
hidden_size = 64
batch_size = 512
n_hidden_layers = 1
ds = sklearn.datasets.load_digits(n_class=10, return_X_y=False)
print(ds.images.shape, ds.images.dtype)
images = tf.reshape(tf.cast(ds.images, dtype), (-1, 64))
labels = tf.cast(ds.target, tf.int64)  # tf.one_hot(ds.target, 10)
def func(x, l, params, get_acc=False):
    # TODO. tensorflow doesnt have a lowdim matmul available in the python api...
    # https://github.com/google/gemmlowp
    h = tf.matmul(x, params['weights_first']) + params['biases_first']
    for i in range(n_hidden_layers):
        h = tf.matmul(h, params['weights_{}'.format(i)]) + params['biases_{}'.format(i)]
    y = tf.matmul(h, params['weights_last']) + params['biases_last']

    p = tf.nn.softmax(y)
    if get_acc:
        return tf.reduce_mean(tf.cast(tf.equal(tf.argmax(p, axis=1), l), 
                             tf.float32))
    else:
        return tf.losses.sparse_softmax_cross_entropy(labels=l,logits=y)
    
params = {'weights_first': tf.random_normal((64, hidden_size), stddev=np.sqrt(2./64), dtype=dtype), 
          'biases_first': tf.random_normal((1, hidden_size), stddev=np.sqrt(2./64), dtype=dtype),
          'weights_last': tf.random_normal((hidden_size, 10), stddev=np.sqrt(2./hidden_size), dtype=dtype), 
          'biases_last': tf.random_normal((1, 10), stddev=np.sqrt(2./hidden_size), dtype=dtype)}
weights = {'weights_{}'.format(i): tf.random_normal((hidden_size, hidden_size), 
                                                    stddev=np.sqrt(2./hidden_size), 
                                                    dtype=dtype)
          for i in range(n_hidden_layers)}
biases = {'biases_{}'.format(i): tf.random_normal((1, hidden_size), 
                                                    stddev=np.sqrt(2./hidden_size), 
                                                    dtype=dtype)
         for i in range(n_hidden_layers)}
params = dict(params, **weights)
params = dict(params, **biases)

# def create(shape, dtype):
#     return tf.cast(tf.random_uniform(shape, minval=-2**7, maxval=2**7), dtype)

# params = {'weights1':create((64, hidden_size), dtype), 
#           'biases1': create((1, hidden_size), dtype),
#           'weights2': create((hidden_size, 10), dtype), 
#           'biases2': create((1, 10), dtype)}

o = func(images, labels, params)
def project(M, v):
    shape = M.get_shape()
    M = tf.reshape(M, (tf.shape(M)[0], -1))
    sim = tf.losses.cosine_distance(tf.nn.l2_normalize(x, 0), tf.nn.l2_normalize(y, 0), dim=0)
    M *= sim
    return tf.reshape(M, shape)

def create_noise(params, N, momentum, stddev):
    return {k: tf.random_normal((N, v.shape[0], v.shape[1]), 
                                dtype=dtype, 
                                stddev=stddev)
           for k, v in params.items()}

def mutate_and_evaluate(func, params, images, labels, n_pop, n_iters, step, momentum):
    noise = create_noise(params, n_pop, momentum, stddev=0.1*(1-step/n_iters))
    losses = []

    # for each elem in the population
    for i in range(n_pop):  # this should be parallelised!
        current_params = {k: params[k]+v[i, ...]
                         for k, v in noise.items()}

        # do we really need to evaluate on all the data?
        idx = np.random.randint(0, images.shape[0]-batch_size)
        # if I am going to do this I really need some mometum
        # as this algol has no memory. so good results are forgotten

        # use our fittness fn/loss fn to evaluate
        loss = func(images[idx: idx+batch_size, ...], 
                    labels[idx: idx+batch_size, ...], 
                    current_params)
        losses.append(loss)
    #     print({k: v.shape for k,v in current_params.items()})

    losses = tf.reshape(tf.stack(losses), (1, n_pop))
    return noise, losses

def update_params(losses, noise, params, 
                  lr=1.0, weighted_sum=True):
    new_params = {}
    m, v = tf.nn.moments(losses, 1)
    losses = (losses - m)/v

    momentum = {}
    for k in params.keys():
        # take a weighted sum
        if weighted_sum:
            dp = tf.reshape(tf.matmul(losses, 
                                      tf.reshape(noise[k], 
                                                (n_pop, -1))),
                            params[k].shape)
        # or could just pick the best. argmax.
        else:
            # this one seems to work better!?
            dp = -noise[k][tf.argmin(losses, axis=1)[0]]
            # elif:
            # would also like to be able to sub in gradients from AD
        
        momentum[k] = dp

#         print(np.sum(np.abs(dp)))
        # do gradient descent on the params
        new_params[k] = params[k] - lr*dp

    return new_params, momentum
class Optimiser():
    def __init__(self):
        pass

class ES(Optimiser):
    pass

class SGD(Optimiser):
    pass
    
iters=1000
momentum = {k: tf.zeros_like(v) for k,v in params.items()}
for i in range(iters):
    noise, losses = mutate_and_evaluate(func, params, images, labels, n_pop, iters, i, momentum)
    params, momentum = update_params(losses, noise, params)
#     print(np.mean(losses))
    acc = func(images, labels, params, get_acc=True)
    print('\rstep: {} loss: {}'.format(i, acc), end='', flush=True)
def train(method, N):
    if method is 'es':
        pass
    if method is 'sgd':
        pass
    
    
es_loss = train('es', 3)
sgd_loss = train('sgd', 3)

plot(es_loss)
plot(sgd_loss)
# want to see that; 
# - es is slow
# - sgd achieves lower accuracy (on linear nets)
# - momentum!?<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13127908842882363, 'fraction_numerical': 0.013480982185844969, 'mean_word_length': 2.4166666666666665, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 11, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '195133', 'license': 'MIT', 'n_tokens_mistral': 2080, 'n_tokens_neox': 2010, 'n_words': 536, 'path': 'gradient-estimation/nonlin-linearNN-withES.ipynb', 'repo_name': 'act65/datasci', 'size': '6231'}