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)