docker-python/tests/test_flax.py at tmp-gpu-test · Kaggle/docker-python · GitHub

History

52 lines (41 loc) · 1.66 KB

Raw

import unittest

import jax

import jax.numpy as jnp

import numpy as np

import optax

from flax import linen as nn

from flax.training import train_state

class TestFlax(unittest.TestCase):

def test_pooling(self):

x = jnp.full((1, 3, 3, 1), 2.)

mul_reduce = lambda x, y: x * y

y = nn.pooling.pool(x, 1., mul_reduce, (2, 2), (1, 1), 'VALID')

np.testing.assert_allclose(y, np.full((1, 2, 2, 1), 2. ** 4))

def test_cnn(self):

class CNN(nn.Module):

@nn.compact

def __call__(self, x):

x = nn.Conv(features=32, kernel_size=(3, 3))(x)

x = nn.relu(x)

x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))

x = nn.Conv(features=64, kernel_size=(3, 3))(x)

x = nn.relu(x)

x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))

x = x.reshape((x.shape[0], -1))

x = nn.Dense(features=256)(x)

x = nn.relu(x)

x = nn.Dense(features=120)(x)

x = nn.log_softmax(x)

return x

def create_train_state(rng, learning_rate, momentum):

cnn = CNN()

params = cnn.init(rng, jnp.ones([1, 224, 224, 3]))['params']

tx = optax.sgd(learning_rate, momentum)

return train_state.TrainState.create(

apply_fn=cnn.apply, params=params, tx=tx)

rng = jax.random.PRNGKey(0)

rng, init_rng = jax.random.split(rng)

learning_rate = 2e-5

momentum = 0.9

state = create_train_state(init_rng, learning_rate, momentum)

self.assertEqual(0, state.step)

docker-python/tests/test_flax.py at tmp-gpu-test · Kaggle/docker-python · GitHub

Expand file tree

Breadcrumbs

test_flax.py

Latest commit

History

Breadcrumbs

test_flax.py

File metadata and controls