Method test_cnn

tests/test_flax.py:20–50

(self)

        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)