| 80 | return seq_matrix, Y |
| 81 | |
| 82 | def DeepSTARR(params): |
| 83 | if params['encode'] == 'one-hot': |
| 84 | input = kl.Input(shape=(249, 4)) |
| 85 | elif params['encode'] == 'k-mer': |
| 86 | input = kl.Input(shape=(1, 64)) |
| 87 | |
| 88 | for i in range(params['convolution_layers']['n_layers']): |
| 89 | x = kl.Conv1D(params['convolution_layers']['filters'][i], |
| 90 | kernel_size = params['convolution_layers']['kernel_sizes'][i], |
| 91 | padding = params['pad'], |
| 92 | name=str('Conv1D_'+str(i+1)))(input) |
| 93 | x = kl.BatchNormalization()(x) |
| 94 | x = kl.Activation('relu')(x) |
| 95 | if params['encode'] == 'one-hot': |
| 96 | x = kl.MaxPooling1D(2)(x) |
| 97 | |
| 98 | if params['dropout_conv'] == 'yes': x = kl.Dropout(params['dropout_prob'])(x) |
| 99 | |
| 100 | # optional attention layers |
| 101 | for i in range(params['transformer_layers']['n_layers']): |
| 102 | if i == 0: |
| 103 | x = x + keras_nlp.layers.SinePositionEncoding()(x) |
| 104 | x = TransformerEncoder(intermediate_dim = params['transformer_layers']['attn_key_dim'][i], |
| 105 | num_heads = params['transformer_layers']['attn_heads'][i], |
| 106 | dropout = params['dropout_prob'])(x) |
| 107 | |
| 108 | # After the convolutional layers, the output is flattened and passed through a series of fully connected/dense layers |
| 109 | # Flattening converts a multi-dimensional input (from the convolutions) into a one-dimensional array (to be connected with the fully connected layers |
| 110 | x = kl.Flatten()(x) |
| 111 | |
| 112 | # Fully connected layers |
| 113 | # Each fully connected layer is followed by batch normalization, ReLU activation, and dropout |
| 114 | for i in range(params['n_dense_layer']): |
| 115 | x = kl.Dense(params['dense_neurons'+str(i+1)], |
| 116 | name=str('Dense_'+str(i+1)))(x) |
| 117 | x = kl.BatchNormalization()(x) |
| 118 | x = kl.Activation('relu')(x) |
| 119 | x = kl.Dropout(params['dropout_prob'])(x) |
| 120 | |
| 121 | # Main model bottleneck |
| 122 | bottleneck = x |
| 123 | |
| 124 | # heads per task (developmental and housekeeping enhancer activities) |
| 125 | # The final output layer is a pair of dense layers, one for each task (developmental and housekeeping enhancer activities), each with a single neuron and a linear activation function |
| 126 | tasks = ['Dev', 'Hk'] |
| 127 | outputs = [] |
| 128 | for task in tasks: |
| 129 | outputs.append(kl.Dense(1, activation='linear', name=str('Dense_' + task))(bottleneck)) |
| 130 | |
| 131 | # Build Keras model object |
| 132 | model = Model([input], outputs) |
| 133 | model.compile(Adam(learning_rate=params['lr']), # Adam optimizer |
| 134 | loss=['mse', 'mse'], # loss is Mean Squared Error (MSE) |
| 135 | loss_weights=[1, 1]) # in case we want to change the weights of each output. For now keep them with same weights |
| 136 | |
| 137 | return model, params |
| 138 | |
| 139 | def train(selected_model, X_train, Y_train, X_valid, Y_valid, params): |