def test_cosine_similarity(): from ntm.similarities import cosine_similarity key_var, memory_var = T.tensor3s('key', 'memory') cosine_similarity_fn = theano.function([key_var, memory_var], \ cosine_similarity(key_var, memory_var, eps=1e-6)) test_key = np.random.rand(16, 4, 20) test_memory = np.random.rand(16, 128, 20) test_output = cosine_similarity_fn(test_key, test_memory) test_output_manual = np.zeros_like(test_output) for i in range(16): for j in range(4): for k in range(128): test_output_manual[i, j, k] = np.dot(test_key[i, j], test_memory[i, k]) / \ np.sqrt(np.sum(test_key[i, j] * test_key[i, j]) * np.sum(test_memory[i, k] * \ test_memory[i, k]) + 1e-6) assert np.allclose(test_output, test_output_manual)
def test_content_addressing(): from ntm.similarities import cosine_similarity beta_var, key_var, memory_var = T.tensor3s('beta', 'key', 'memory') beta_var = T.addbroadcast(beta_var, 2) betaK = beta_var * cosine_similarity(key_var, memory_var) w_c = lasagne.nonlinearities.softmax(betaK.reshape((16 * 4, 128))) w_c = w_c.reshape(betaK.shape) content_addressing_fn = theano.function([beta_var, key_var, memory_var], w_c) beta = np.random.rand(16, 4, 1) key = np.random.rand(16, 4, 20) memory = np.random.rand(16, 128, 20) weights = content_addressing_fn(beta, key, memory) weights_manual = np.zeros_like(weights) def softmax(x): y = np.exp(x.T - np.max(x, axis=1)) z = y / np.sum(y, axis=0) return z.T betaK_manual = np.zeros((16, 4, 128)) for i in range(16): for j in range(4): for k in range(128): betaK_manual[i, j, k] = beta[i, j, 0] * np.dot(key[i, j], \ memory[i, k]) / np.sqrt(np.sum(key[i, j] * key[i, j]) * \ np.sum(memory[i, k] * memory[i, k]) + 1e-6) for i in range(16): weights_manual[i] = softmax(betaK_manual[i]) assert weights.shape == (16, 4, 128) assert np.allclose(np.sum(weights, axis=2), np.ones((16, 4))) assert np.allclose(weights, weights_manual)