Function test_NCELoss

numpy_ml/tests/test_nn.py:241–331

(N=1)

def test_NCELoss(N=1):
    from numpy_ml.neural_nets.losses import NCELoss
    from numpy_ml.utils.data_structures import DiscreteSampler

    np.random.seed(12345)

    N = np.inf if N is None else N

    i = 1
    while i < N + 1:
        n_ex = np.random.randint(1, 10)
        n_c = np.random.randint(1, 10)
        n_out = np.random.randint(1, 300)
        vocab_size = np.random.randint(200, 1000)
        num_negative_samples = np.random.randint(1, 10)

        embeddings = random_tensor((n_ex, n_c, n_out), standardize=True)
        target = np.random.randint(0, vocab_size, (n_ex, 1))

        probs = np.random.rand(vocab_size)
        probs /= probs.sum()

        D = DiscreteSampler(probs, log=False, with_replacement=False)
        NCE = NCELoss(vocab_size, D, num_negative_samples)
        my_loss, _ = NCE(embeddings, target.flatten())

        my_dLdX = NCE.grad(update_params=False)
        my_dLdW = NCE.gradients["W"]
        my_dLdb = NCE.gradients["b"]

        NCE.gradients["W"] = np.zeros_like(NCE.parameters["W"])
        NCE.gradients["b"] = np.zeros_like(NCE.parameters["b"])

        MY_final_loss, TF_final_loss = 0, 0
        MY_dLdX, TF_dLdX = np.zeros_like(embeddings), np.zeros_like(embeddings)
        TF_dLdW, TF_dLdb = (
            np.zeros_like(NCE.parameters["W"]),
            np.zeros_like(NCE.parameters["b"]),
        )

        # XXX: instead of calculating the tf NCE on the entire batch, we
        # calculate it per-example and then sum. this is really lame and should
        # be changed to operate on batches.
        nv = NCE.derived_variables["noise_samples"][0]
        for ix, emb in enumerate(embeddings):
            sv = (nv[0], np.array([nv[1][0, ix]]), nv[2])

            NCE.X = []
            for k, v in NCE.derived_variables.items():
                NCE.derived_variables[k] = []

            for k, v in NCE.gradients.items():
                NCE.gradients[k] = np.zeros_like(v)

            my, _ = NCE(emb[None, :, :], target[ix], neg_samples=sv[0])

            NCE.derived_variables["noise_samples"] = [sv]
            dldx = NCE.grad(update_params=False)