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_)