def identity_block(inputs):
main = Sequential(Conv(filters1, (1, 1)), BatchNorm(), relu,
Conv(filters2, (ks, ks), padding='SAME'),
BatchNorm(), relu, Conv(inputs.shape[3], (1, 1)),
BatchNorm())
return relu(sum((main(inputs), inputs)))
def conv_block(inputs):
main = Sequential(Conv(filters1, (1, 1), strides), BatchNorm(), relu,
Conv(filters2, (ks, ks),
padding='SAME'), BatchNorm(), relu,
Conv(filters3, (1, 1)), BatchNorm())
shortcut = Sequential(Conv(filters3, (1, 1), strides), BatchNorm())
return relu(sum((main(inputs), shortcut(inputs))))
def test_regularized_submodule():
net = Sequential(Conv(2, (1, 1)), relu, Conv(2, (1, 1)), relu, flatten,
L2Regularized(Sequential(Dense(2), relu, Dense(2), np.sum), .1))
input = np.ones((1, 3, 3, 1))
params = net.init_parameters(input, key=PRNGKey(0))
assert (2, 2) == params.regularized.model.dense1.kernel.shape
out = net.apply(params, input)
assert () == out.shape
def test_Conv_runs(channels, filter_shape, padding, strides, input_shape,
dilation):
conv = Conv(channels,
filter_shape,
strides=strides,
padding=padding,
dilation=dilation)
inputs = random_inputs(input_shape)
params = conv.init_parameters(PRNGKey(0), inputs)
conv.apply(params, inputs)
def test_reparametrized_submodule():
net = Sequential(
Conv(2, (3, 3)), relu, Conv(2, (3, 3)), relu, flatten,
Reparametrized(Sequential(Dense(2), relu, Dense(2)), Scaled))
input = np.ones((1, 3, 3, 1))
params = net.init_parameters(PRNGKey(0), input)
assert (2, 2) == params.reparametrized.model.dense1.kernel.shape
out = net.apply(params, input)
assert (1, 2) == out.shape
def test_flatten_shape():
conv = Conv(2,
filter_shape=(3, 3),
padding='SAME',
kernel_init=zeros,
bias_init=zeros)
inputs = np.zeros((1, 5, 5, 2))
params = conv.init_parameters(PRNGKey(0), inputs)
out = conv.apply(params, inputs)
assert np.array_equal(np.zeros((1, 5, 5, 2)), out)
flattened = Sequential(conv, flatten)
out = flattened.apply_from({conv: params}, inputs)
assert np.array_equal(np.zeros((1, 50)), out)
def test_pool_shape(Pool):
conv = Conv(2,
filter_shape=(3, 3),
padding='SAME',
kernel_init=zeros,
bias_init=zeros)
inputs = np.zeros((1, 5, 5, 2))
pooled = Sequential(conv, Pool(window_shape=(1, 1), strides=(2, 2)))
params = pooled.init_parameters(PRNGKey(0), inputs)
out = pooled.apply(params, inputs)
assert np.array_equal(np.zeros((1, 3, 3, 2)), out)
def test_external_sequential_submodule():
layer = Sequential(Conv(4, (2, 2)), flatten, relu, Dense(3), relu,
Dense(2), Sequential(Dense(2), relu))
inputs = np.zeros((1, 5, 5, 2))
params = layer.init_parameters(inputs, key=PRNGKey(0))
assert (4, ) == params.conv.bias.shape
assert (3, ) == params.dense0.bias.shape
assert (3, 2) == params.dense1.kernel.shape
assert (2, ) == params.dense1.bias.shape
assert (2, ) == params.sequential.dense.bias.shape
out = layer.apply(params, inputs)
assert (1, 2) == out.shape
out_ = layer.apply(params, inputs, jit=True)
assert np.allclose(out, out_)