| 19 | |
| 20 | def test_cnn(self): |
| 21 | class CNN(nn.Module): |
| 22 | @nn.compact |
| 23 | def __call__(self, x): |
| 24 | x = nn.Conv(features=32, kernel_size=(3, 3))(x) |
| 25 | x = nn.relu(x) |
| 26 | x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2)) |
| 27 | x = nn.Conv(features=64, kernel_size=(3, 3))(x) |
| 28 | x = nn.relu(x) |
| 29 | x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2)) |
| 30 | x = x.reshape((x.shape[0], -1)) |
| 31 | x = nn.Dense(features=256)(x) |
| 32 | x = nn.relu(x) |
| 33 | x = nn.Dense(features=120)(x) |
| 34 | x = nn.log_softmax(x) |
| 35 | return x |
| 36 | |
| 37 | def create_train_state(rng, learning_rate, momentum): |
| 38 | cnn = CNN() |
no outgoing calls