Mercurial > repos > bgruening > create_tool_recommendation_model

comparison train_transformer.py @ 6:e94dc7945639 draft default tip

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload for repository https://github.com/bgruening/galaxytools/tree/recommendation_training/tools/tool_recommendation_model commit 24bab7a797f53fe4bcc668b18ee0326625486164
author bgruening
date Sun, 16 Oct 2022 11:52:10 +0000
parents
children
comparison
equal deleted inserted replaced
5:4f7e6612906b 6:e94dc7945639
1 import tensorflow as tf
2 import transformer_network
3 import utils
4 from tensorflow.keras.layers import (Dense, Dropout, GlobalAveragePooling1D,
5 Input)
6 from tensorflow.keras.models import Model
7
8
9 def create_model(vocab_size, config):
10 embed_dim = config["embedding_dim"]
11 ff_dim = config["feed_forward_dim"]
12 max_len = config["maximum_path_length"]
13 dropout = config["dropout"]
14
15 inputs = Input(shape=(max_len,))
16 embedding_layer = transformer_network.TokenAndPositionEmbedding(max_len, vocab_size, embed_dim)
17 x = embedding_layer(inputs)
18 transformer_block = transformer_network.TransformerBlock(embed_dim, config["n_heads"], ff_dim)
19 x, weights = transformer_block(x)
20 x = GlobalAveragePooling1D()(x)
21 x = Dropout(dropout)(x)
22 x = Dense(ff_dim, activation="relu")(x)
23 x = Dropout(dropout)(x)
24 outputs = Dense(vocab_size, activation="sigmoid")(x)
25 return Model(inputs=inputs, outputs=[outputs, weights])
26
27
28 def create_enc_transformer(train_data, train_labels, test_data, test_labels, f_dict, r_dict, c_wts, c_tools, pub_conn, tr_t_freq, config):
29 print("Train transformer...")
30 vocab_size = len(f_dict) + 1
31
32 enc_optimizer = tf.keras.optimizers.Adam(learning_rate=config["learning_rate"])
33
34 model = create_model(vocab_size, config)
35
36 u_tr_y_labels, u_tr_y_labels_dict = utils.get_u_tr_labels(train_labels)
37 u_te_y_labels, u_te_y_labels_dict = utils.get_u_tr_labels(test_labels)
38
39 trained_on_labels = [int(item) for item in list(u_tr_y_labels_dict.keys())]
40
41 epo_tr_batch_loss = list()
42 epo_tr_batch_acc = list()
43 all_sel_tool_ids = list()
44
45 te_lowest_t_ids = utils.get_low_freq_te_samples(test_data, test_labels, tr_t_freq)
46 tr_log_step = config["tr_logging_step"]
47 te_log_step = config["te_logging_step"]
48 n_train_steps = config["n_train_iter"]
49 te_batch_size = config["te_batch_size"]
50 tr_batch_size = config["tr_batch_size"]
51 sel_tools = list()
52 for batch in range(n_train_steps):
53 x_train, y_train, sel_tools = utils.sample_balanced_tr_y(train_data, train_labels, u_tr_y_labels_dict, tr_batch_size, tr_t_freq, sel_tools)
54 all_sel_tool_ids.extend(sel_tools)
55 with tf.GradientTape() as model_tape:
56 prediction, att_weights = model(x_train, training=True)
57 tr_loss, tr_cat_loss = utils.compute_loss(y_train, prediction)
58 tr_acc = tf.reduce_mean(utils.compute_acc(y_train, prediction))
59 trainable_vars = model.trainable_variables
60 model_gradients = model_tape.gradient(tr_loss, trainable_vars)
61 enc_optimizer.apply_gradients(zip(model_gradients, trainable_vars))
62 epo_tr_batch_loss.append(tr_loss.numpy())
63 epo_tr_batch_acc.append(tr_acc.numpy())
64 if (batch + 1) % tr_log_step == 0:
65 print("Total train data size: ", train_data.shape, train_labels.shape)
66 print("Batch train data size: ", x_train.shape, y_train.shape)
67 print("At Step {}/{} training loss:".format(str(batch + 1), str(n_train_steps)))
68 print(tr_loss.numpy())
69 if (batch + 1) % te_log_step == 0:
70 print("Predicting on test data...")
71 utils.validate_model(test_data, test_labels, te_batch_size, model, f_dict, r_dict, u_te_y_labels_dict, trained_on_labels, te_lowest_t_ids)
72 print("Saving model after training for {} steps".format(n_train_steps))
73 utils.save_model_file(model, r_dict, c_wts, c_tools, pub_conn, config["trained_model_path"])