134 lines
3.5 KiB
Cython
134 lines
3.5 KiB
Cython
cimport numpy as np
|
|
import numpy as np
|
|
|
|
from sys import stderr
|
|
|
|
from libc.stdlib cimport malloc
|
|
from libc.string cimport memcpy
|
|
|
|
import nn
|
|
|
|
|
|
X_train, y_train, X_test, y_test = nn.load_mnist()
|
|
|
|
|
|
cdef extern from "numpy/arrayobject.h":
|
|
void *PyArray_DATA(np.ndarray arr)
|
|
|
|
|
|
ctypedef public struct Weight:
|
|
size_t dims
|
|
long* shape
|
|
float* W
|
|
|
|
|
|
ctypedef public struct WeightList:
|
|
size_t n_weights;
|
|
Weight* weights;
|
|
|
|
|
|
cdef public char *greeting():
|
|
return f'The value is {3**3**3}'.encode('utf-8')
|
|
|
|
|
|
cdef public void debug_print(object o):
|
|
print(o)
|
|
|
|
|
|
cdef public object create_network():
|
|
return nn.create_mnist_network()
|
|
|
|
|
|
cdef public void set_net_weights(object net, WeightList* wl):
|
|
net.set_weights(wrap_weight_list(wl))
|
|
|
|
|
|
cdef public void step_net(
|
|
object net, float* X, float* y, size_t batch_size
|
|
):
|
|
in_shape = (batch_size,) + net.layers[0].input_shape[1:]
|
|
out_shape = (batch_size,) + net.layers[-1].output_shape[1:]
|
|
X_train = np.asarray(<float[:np.prod(in_shape)]>X).reshape(in_shape)
|
|
y_train = np.asarray(<float[:np.prod(out_shape)]>y).reshape(out_shape)
|
|
|
|
net.train_on_batch(X_train, y_train)
|
|
|
|
|
|
cdef public float eval_net(object net):
|
|
return net.evaluate(X_test, y_test, verbose=False)[1]
|
|
|
|
|
|
cdef public void mnist_batch(float* X, float* y, size_t bs,
|
|
int part, int total):
|
|
if total == 0:
|
|
X_pool, y_pool = X_train, y_train
|
|
else:
|
|
partsize = len(X_train) // total
|
|
X_pool = X_train[part*partsize:(part+1)*partsize]
|
|
y_pool = y_train[part*partsize:(part+1)*partsize]
|
|
|
|
idx = np.random.choice(len(X_pool), bs, replace=True)
|
|
|
|
X_r = X_pool[idx]
|
|
y_r = y_pool[idx]
|
|
|
|
assert X_r.flags['C_CONTIGUOUS']
|
|
assert y_r.flags['C_CONTIGUOUS']
|
|
memcpy(X, <float*>PyArray_DATA(X_r), X_r.size * sizeof(float))
|
|
memcpy(y, <float*>PyArray_DATA(y_r), y_r.size * sizeof(float))
|
|
|
|
|
|
cdef public void init_weightlist_like(WeightList* wl, object net):
|
|
weights = net.get_weights()
|
|
wl.n_weights = len(weights)
|
|
wl.weights = <Weight*>malloc(sizeof(Weight) * wl.n_weights)
|
|
for i, w in enumerate(weights):
|
|
sh = np.asarray(w.shape, dtype=long)
|
|
wl.weights[i].dims = sh.size
|
|
wl.weights[i].shape = <long*>malloc(sizeof(long) * sh.size)
|
|
wl.weights[i].W = <float*>malloc(sizeof(float) * w.size)
|
|
|
|
assert sh.flags['C_CONTIGUOUS']
|
|
memcpy(wl.weights[i].shape, <long*>PyArray_DATA(sh),
|
|
sh.size * sizeof(long))
|
|
|
|
|
|
cdef public void update_weightlist(WeightList* wl, object net):
|
|
weights = net.get_weights()
|
|
for i, w in enumerate(weights):
|
|
w = w.astype(np.float32)
|
|
|
|
assert w.flags['C_CONTIGUOUS']
|
|
memcpy(wl.weights[i].W, <float*>PyArray_DATA(w),
|
|
w.size * sizeof(float))
|
|
|
|
|
|
cdef public void combo_weights(
|
|
WeightList* wl_frank, WeightList* wls, size_t num_weights
|
|
):
|
|
"""Not a one-liner anymore :/"""
|
|
alpha = 1. / num_weights
|
|
frank = wrap_weight_list(wl_frank)
|
|
for w in frank:
|
|
w[:] = 0
|
|
for i in range(num_weights):
|
|
for wf, ww in zip(frank, wrap_weight_list(&wls[i])):
|
|
wf += alpha * ww
|
|
|
|
|
|
cdef list wrap_weight_list(WeightList* wl):
|
|
weights = []
|
|
for i in range(wl.n_weights):
|
|
w_shape = <long[:wl.weights[i].dims]>wl.weights[i].shape
|
|
w_numel = np.prod(w_shape)
|
|
weights.append(
|
|
np.asarray(<float[:w_numel]>wl.weights[i].W).reshape(w_shape)
|
|
)
|
|
return weights
|
|
|
|
|
|
def inspect_array(a):
|
|
print(a.flags, flush=True)
|
|
print(a.dtype, flush=True)
|
|
print(a.sum(), flush=True)
|