Function test_WaveNetModule

numpy_ml/tests/test_nn.py:2072–2166

(N=10)

def test_WaveNetModule(N=10):
    from numpy_ml.neural_nets.modules import WavenetResidualModule

    N = np.inf if N is None else N

    np.random.seed(12345)

    i = 1
    while i < N + 1:
        n_ex = np.random.randint(1, 10)
        l_in = np.random.randint(1, 10)
        ch_residual, ch_dilation = np.random.randint(1, 5), np.random.randint(1, 5)
        f_width = min(l_in, np.random.randint(1, 5))
        d = np.random.randint(0, 5)

        X_main = np.zeros_like(
            random_tensor((n_ex, l_in, ch_residual), standardize=True)
        )
        X_main[0][0][0] = 1.0
        X_skip = np.zeros_like(
            random_tensor((n_ex, l_in, ch_residual), standardize=True)
        )

        # initialize Conv2D layer
        L1 = WavenetResidualModule(
            ch_residual=ch_residual,
            ch_dilation=ch_dilation,
            kernel_width=f_width,
            dilation=d,
        )

        # forward prop
        Y_main, Y_skip = L1.forward(X_main, X_skip)

        # backprop
        dLdY_skip = np.ones_like(Y_skip)
        dLdY_main = np.ones_like(Y_main)
        dLdX_main, dLdX_skip = L1.backward(dLdY_skip, dLdY_main)

        _, conv_1x1_pad = pad1D(
            L1._dv["multiply_gate_out"], "same", kernel_width=1, stride=1, dilation=0
        )
        if conv_1x1_pad[0] != conv_1x1_pad[1]:
            print("Skipping")
            continue

        conv_1x1_pad = conv_1x1_pad[0]

        # get gold standard gradients
        gold_mod = TorchWavenetModule(L1.parameters, L1.hyperparameters, conv_1x1_pad)
        golds = gold_mod.extract_grads(X_main, X_skip)

        dv = L1.derived_variables
        pc = L1.parameters["components"]
        gr = L1.gradients["components"]

        params = [
            (L1.X_main, "X_main"),