| 30 | |
| 31 | |
| 32 | class TNN: |
| 33 | def __init__(self, V, D, K, activation): |
| 34 | self.D = D |
| 35 | self.f = activation |
| 36 | |
| 37 | # word embedding |
| 38 | We = init_weight(V, D) |
| 39 | |
| 40 | # linear terms |
| 41 | W1 = init_weight(D, D) |
| 42 | W2 = init_weight(D, D) |
| 43 | |
| 44 | # bias |
| 45 | bh = np.zeros(D) |
| 46 | |
| 47 | # output layer |
| 48 | Wo = init_weight(D, K) |
| 49 | bo = np.zeros(K) |
| 50 | |
| 51 | # make them tensorflow variables |
| 52 | self.We = tf.Variable(We.astype(np.float32)) |
| 53 | self.W1 = tf.Variable(W1.astype(np.float32)) |
| 54 | self.W2 = tf.Variable(W2.astype(np.float32)) |
| 55 | self.bh = tf.Variable(bh.astype(np.float32)) |
| 56 | self.Wo = tf.Variable(Wo.astype(np.float32)) |
| 57 | self.bo = tf.Variable(bo.astype(np.float32)) |
| 58 | self.params = [self.We, self.W1, self.W2, self.Wo] |
| 59 | |
| 60 | def fit(self, trees, lr=1e-1, mu=0.9, reg=0.1, epochs=5): |
| 61 | train_ops = [] |
| 62 | costs = [] |
| 63 | predictions = [] |
| 64 | all_labels = [] |
| 65 | i = 0 |
| 66 | N = len(trees) |
| 67 | print("Compiling ops") |
| 68 | for t in trees: |
| 69 | i += 1 |
| 70 | sys.stdout.write("%d/%d\r" % (i, N)) |
| 71 | sys.stdout.flush() |
| 72 | logits = self.get_output(t) |
| 73 | labels = get_labels(t) |
| 74 | all_labels.append(labels) |
| 75 | |
| 76 | cost = self.get_cost(logits, labels, reg) |
| 77 | costs.append(cost) |
| 78 | |
| 79 | prediction = tf.argmax(input=logits, axis=1) |
| 80 | predictions.append(prediction) |
| 81 | |
| 82 | train_op = tf.compat.v1.train.MomentumOptimizer(lr, mu).minimize(cost) |
| 83 | train_ops.append(train_op) |
| 84 | |
| 85 | # save for later so we don't have to recompile |
| 86 | self.predictions = predictions |
| 87 | self.all_labels = all_labels |
| 88 | self.saver = tf.compat.v1.train.Saver() |
| 89 | |