def test_submodule_reuse():
inputs = np.zeros((1, 2))
layer = Dense(5)
net1 = Sequential(layer, Dense(2))
net2 = Sequential(layer, Dense(3))
layer_params = layer.init_parameters(inputs, key=PRNGKey(0))
net1_params = net1.init_parameters(inputs,
key=PRNGKey(1),
reuse={layer: layer_params})
net2_params = net2.init_parameters(inputs,
key=PRNGKey(2),
reuse={layer: layer_params})
out1 = net1.apply(net1_params, inputs)
assert out1.shape == (1, 2)
out2 = net2.apply(net2_params, inputs)
assert out2.shape == (1, 3)
assert_dense_parameters_equal(layer_params, net1_params[0])
assert_dense_parameters_equal(layer_params, net2_params[0])
new_layer_params = layer.init_parameters(inputs, key=PRNGKey(3))
combined_params = net1.parameters_from(
{
net1: net1_params,
layer: new_layer_params
}, inputs)
assert_dense_parameters_equal(new_layer_params, combined_params.dense0)
assert_dense_parameters_equal(net1_params.dense1, combined_params.dense1)
def test_no_reuse():
inputs = np.zeros((1, 2))
layer = Dense(5)
net1 = Sequential(layer, Dense(2))
p1 = net1.init_parameters(inputs, key=PRNGKey(0))
net2 = Sequential(layer, Dense(3))
p2 = net2.init_parameters(inputs, key=PRNGKey(1))
assert p1[0].kernel.shape == p2[0].kernel.shape
assert p1[0].bias.shape == p2[0].bias.shape
assert not np.array_equal(p1[0][0], p2[0][0])
assert not np.array_equal(p1[0][1], p2[0][1])
def test_ocr_rnn():
length = 5
carry_size = 3
class_count = 4
inputs = np.zeros((1, length, 4))
def rnn():
return Rnn(*GRUCell(carry_size, zeros))
net = Sequential(
rnn(),
rnn(),
rnn(),
lambda x: np.reshape(x, (-1, carry_size)
), # -> same weights for all time steps
Dense(class_count, zeros, zeros),
softmax,
lambda x: np.reshape(x, (-1, length, class_count)))
params = net.init_parameters(PRNGKey(0), inputs)
assert len(params) == 4
cell = params.rnn0.gru_cell
assert len(cell) == 3
assert np.array_equal(np.zeros((7, 3)), cell.update_kernel)
assert np.array_equal(np.zeros((7, 3)), cell.reset_kernel)
assert np.array_equal(np.zeros((7, 3)), cell.compute_kernel)
out = net.apply(params, inputs)
assert np.array_equal(.25 * np.ones((1, 5, 4)), out)
def test_submodule_reuse():
inputs = np.zeros((1, 2))
layer = Dense(5)
net1 = Sequential(layer, Dense(2))
net2 = Sequential(layer, Dense(3))
layer_params = layer.init_parameters(PRNGKey(0), inputs)
net1_params = net1.init_parameters(PRNGKey(1), inputs, reuse={layer: layer_params})
net2_params = net2.init_parameters(PRNGKey(2), inputs, reuse={layer: layer_params})
out1 = net1.apply(net1_params, inputs)
assert out1.shape == (1, 2)
out2 = net2.apply(net2_params, inputs)
assert out2.shape == (1, 3)
assert_dense_params_equal(layer_params, net1_params[0])
assert_dense_params_equal(layer_params, net2_params[0])
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_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_reuse_api():
inputs = np.zeros((1, 2))
net = Dense(5)
net_params = net.init_parameters(inputs, key=PRNGKey(0))
# train net params...
transfer_net = Sequential(net, relu, Dense(2))
transfer_net_params = transfer_net.init_parameters(inputs, key=PRNGKey(1),
reuse={net: net_params})
assert net_params == transfer_net_params.dense0
def test_input_dependent_nested_modules():
@parametrized
def layer(inputs):
return Dense(inputs.shape[0])(inputs)
net = Sequential(Dense(3), layer)
inputs = np.zeros((5, 3))
params = net.init_parameters(inputs, key=PRNGKey(0))
out = net.apply(params, inputs)
assert (5, 5) == out.shape
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_param_sharing():
layer = Dense(2, zeros, zeros)
shared_net = Sequential(layer, layer)
inputs = np.zeros((1, 2))
params = shared_net.init_parameters(inputs, key=PRNGKey(0))
assert_parameters_equal(((np.zeros((2, 2)), np.zeros(2)), ), params)
out = shared_net.apply(params, inputs)
assert np.array_equal(np.zeros((1, 2)), out)
out = shared_net.apply(params, inputs, jit=True)
assert np.array_equal(np.zeros((1, 2)), out)
def test_collection_input(type):
@parametrized
def net(inputs):
assert isinstance(inputs, type)
return inputs[0] * inputs[1] * parameter((), zeros)
inputs = type((np.zeros(2), np.zeros(2)))
params = net.init_parameters(inputs, key=PRNGKey(0))
out = net.apply(params, inputs)
assert np.array_equal(np.zeros(2), out)
net = Sequential(net)
params = net.init_parameters(inputs, key=PRNGKey(0))
out = net.apply(params, inputs)
assert np.array_equal(np.zeros(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_)
def test_parameters_from_shared_submodules():
sublayer = Dense(2)
a = Sequential(sublayer, relu)
b = Sequential(sublayer, np.sum)
@parametrized
def net(inputs):
return a(inputs) * b(inputs)
inputs = np.zeros((1, 3))
a_params = a.init_parameters(inputs, key=PRNGKey(0))
out = a.apply(a_params, inputs)
params = net.parameters_from({a: a_params}, inputs)
assert_parameters_equal(a_params.dense.kernel,
params.sequential0.dense.kernel)
assert_parameters_equal((), params.sequential1)
out = net.apply(params, inputs)
out_ = net.apply_from({a: a_params}, inputs)
assert np.array_equal(out, out_)
out_ = net.apply_from({a: a_params}, inputs, jit=True)
assert np.array_equal(out, out_)
out_ = net.apply_from({a.shaped(inputs): a_params}, inputs)
assert np.array_equal(out, out_)
out_ = net.apply_from({a.shaped(inputs): a_params}, inputs, jit=True)
assert np.array_equal(out, out_)
out_ = net.shaped(inputs).apply_from({a: a_params})
assert np.array_equal(out, out_)
out_ = net.shaped(inputs).apply_from({a: a_params}, jit=True)
assert np.array_equal(out, out_)
out_ = net.shaped(inputs).apply_from({a.shaped(inputs): a_params})
assert np.array_equal(out, out_)
out_ = net.shaped(inputs).apply_from({a.shaped(inputs): a_params},
jit=True)
assert np.array_equal(out, out_)
def test_parameters_from_subsubmodule():
subsublayer = Dense(2)
sublayer = Sequential(subsublayer, relu)
net = Sequential(sublayer, np.sum)
inputs = np.zeros((1, 3))
params = net.init_parameters(inputs, key=PRNGKey(0))
out = net.apply(params, inputs)
subsublayer_params = subsublayer.init_parameters(inputs, key=PRNGKey(0))
params_ = net.parameters_from({subsublayer: subsublayer_params}, inputs)
assert_dense_parameters_equal(subsublayer_params, params_[0][0])
out_ = net.apply(params_, inputs)
assert out.shape == out_.shape
out_ = net.apply_from({subsublayer: subsublayer_params}, inputs)
assert out.shape == out_.shape
out_ = net.apply_from({subsublayer: subsublayer_params}, inputs, jit=True)
assert out.shape == out_.shape
def test_params_from_shared_submodules2():
sublayer = Dense(2)
a = Sequential(sublayer, relu)
b = Sequential(sublayer, np.sum)
@parametrized
def net(inputs):
return a(inputs), b(inputs)
inputs = np.zeros((1, 3))
a_params = a.init_parameters(PRNGKey(0), inputs)
out = a.apply(a_params, inputs)
params = net.parameters_from({a: a_params}, inputs)
assert_dense_params_equal(a_params.dense, params.sequential0.dense)
assert_dense_params_equal(a_params.dense, params.sequential1.dense)
# TODO parameters are duplicated, optimization with weight sharing is wrong:
# TODO instead: assert 1 == len(params)
out_, _ = net.apply(params, inputs)
assert np.array_equal(out, out_)
def test_ocr_rnn():
length = 5
carry_size = 3
class_count = 4
inputs = jnp.zeros((1, length, 4))
def rnn():
return Rnn(*GRUCell(carry_size, zeros))
net = Sequential(
rnn(),
rnn(),
rnn(),
lambda x: jnp.reshape(x, (-1, carry_size)
), # -> same weights for all time steps
Dense(class_count, zeros, zeros),
softmax,
lambda x: jnp.reshape(x, (-1, length, class_count)))
params = net.init_parameters(inputs, key=PRNGKey(0))
assert len(params) == 4
cell = params.rnn0.gru_cell
assert len(cell) == 3
assert jnp.array_equal(jnp.zeros((7, 3)), cell.update_kernel)
assert jnp.array_equal(jnp.zeros((7, 3)), cell.reset_kernel)
assert jnp.array_equal(jnp.zeros((7, 3)), cell.compute_kernel)
out = net.apply(params, inputs)
@parametrized
def cross_entropy(images, targets):
prediction = net(images)
return jnp.mean(-jnp.sum(targets * jnp.log(prediction), (1, 2)))
opt = optimizers.RmsProp(0.003)
state = opt.init(cross_entropy.init_parameters(inputs, out,
key=PRNGKey(0)))
state = opt.update(cross_entropy.apply, state, inputs, out)
opt.update(cross_entropy.apply, state, inputs, out, jit=True)
def test_parametrized_jit_parameter_sharing():
d = Dense(3)
net = Sequential(d, jit(d))
params = net.init_parameters(np.zeros((2, 3)), key=PRNGKey(0))
assert len(params) == 1
net.apply(params, np.zeros((2, 3)))
