def test_atomicnn_default(): n_inputs = 256 default_sizes = [n_inputs, 160, 128, 96, 1] model = models.AtomicNN(n_inputs) assert len(model.layers) == 2 * (len(default_sizes) - 2) + 1 for i, (in_size, out_size) in enumerate(zip(default_sizes[:-1], default_sizes[1:])): assert isinstance(model.layers[2 * i], nn.Linear) assert model.layers[2 * i].weight.shape == (out_size, in_size)
def test_atomicnn_dropout(): n_inputs = 256 sizes = [128, 64, 32, 16, 8, 4, 1] expected_sizes = [n_inputs] + sizes model = models.AtomicNN(n_inputs, sizes, 0.5) assert len(model.layers) == 3 * (len(expected_sizes) - 2) + 1 for i, (in_size, out_size) in enumerate( zip(expected_sizes[:-1], expected_sizes[1:]) ): assert isinstance(model.layers[3 * i], nn.Linear) assert model.layers[3 * i].weight.shape == (out_size, in_size)