Memory usage is too high

Including a Conv into the mnist example

import time

import jax.numpy as np
import numpy.random as npr
from jax.random import PRNGKey

from jaxnet import Sequential, parametrized, Dense, relu, logsoftmax, optimizers, Conv, flatten

def _one_hot(x, k, dtype=np.float32):
    """Create a one-hot encoding of x of size k."""
    return np.array(x[:, None] == np.arange(k), dtype)

def mnist():
    import tensorflow_datasets as tfds
    dataset = tfds.load("mnist:1.0.0")
    images = lambda d: np.reshape(np.float32(d['image']) / 256, (-1, 28, 28, 1))
    labels = lambda d: _one_hot(d['label'], 10)
    train = next(tfds.as_numpy(dataset['train'].shuffle(50000).batch(50000)))
    test = next(tfds.as_numpy(dataset['test'].batch(10000)))
    return images(train), labels(train), images(test), labels(test)

predict = Sequential(
    Conv(32, (5, 5)), relu, flatten,
    Dense(500), relu,
    Dense(10), logsoftmax)

results in out-of-memory on GPU colab during apply_from (init_parameters is fine).

juliuskunze / jaxnet

Memory usage is too high #11