Move to pyTorch form TF!!

Author: Katya Govorkova

.gitignore

@@ -159,4 +159,6 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-output/
+*.DS_Store
+output/
+.snakemake/

Snakefile

@@ -1,13 +1,18 @@
 rule train_model:
-    input:
+    params:
+        data_path = '/home/katya.govorkova/challenge_datasets/ligo_datasets/output'
     output:
-        model = directory('output/saved_model')
+        model = 'output/model.pth'
     shell:
-        'python3 scripts/train.py {output.model}'
+        'python3 scripts/train.py {params.data_path} {output.model}'
+
 
 rule evaluate_on_blackbox:
     input:
-        rules.train_model.output.model
+        model = rules.train_model.output.model
+    params:
+        data_path = '/home/katya.govorkova/challenge_datasets/ligo_datasets/output'
     output:
+        submission = 'output/submission.npy'
     shell:
-        'python3 scripts/evaluate.py'
+        'python3 scripts/evaluate.py {params.data_path} {input.model} {output.submission}'

scripts/dataset.py

@@ -0,0 +1,23 @@
+import torch
+from torch.utils.data import Dataset
+
+
+class TorchADDataset(Dataset):
+  'Characterizes a dataset for PyTorch'
+  def __init__(self, features, labels, device):
+        'Initialization'
+        self.device = device
+        self.features = torch.from_numpy(features).to(dtype=torch.float32, device=self.device)
+        self.labels = torch.from_numpy(labels).to(dtype=torch.float32, device=self.device)
+
+  def __len__(self):
+        'Denotes the total number of samples'
+        return len(self.features)
+
+  def __getitem__(self, index):
+        'Generates one sample of data'
+        # Load data and get label
+        X = self.features[index]
+        y = self.labels[index]
+
+        return X, y

scripts/evaluate.py

@@ -1,11 +1,59 @@
-pretrained_model = tf.keras.models.load_model(os.path.join(data_path, 'saved_model/my_model'))
+import os
+import argparse
+import numpy as np
+from sklearn import metrics
 
-# Check its architecture
-pretrained_model.summary()
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
 
-# load challenge test data
-blackbox = np.load(os.path.join(data_path, 'ligo_blackbox.npz'))['data'].reshape((-1,200,2))
-print('Blackbox shape:', blackbox.shape)
 
-blackbox_prediction = model.predict(blackbox)
-np.save('submission.npy', blackbox_prediction)
+from model import ADClassifier
+
+
+def main(args):
+
+    # Model class must be defined somewhere
+
+    pretrained_model = ADClassifier()
+    pretrained_model.load_state_dict(torch.load(args.modeldir))
+    pretrained_model.eval()
+
+    # Check its architecture
+    print(pretrained_model)
+
+    # load challenge test data
+    blackbox = np.load(os.path.join(args.data_path, 'ligo_blackbox.npz'))['data'].reshape((-1,200,2))
+    print('Blackbox shape:', blackbox.shape)
+
+    # transform to float64
+    x = torch.from_numpy(blackbox)
+    x = x.to(torch.float32)
+
+    blackbox_prediction = pretrained_model(x)
+    np.save(args.submission_path, blackbox_prediction.detach().numpy())
+
+    # scores = pretrained_model(x_val)
+    # fpr, tpr, thresholds = metrics.roc_curve(y_val, scores)
+    # auc = metrics.roc_auc_score(y_val, scores)
+    # print(f'The total AUC is {auc*100:.1f} %')
+    # plt.plot(fpr, tpr)
+    # plt.xlabel('FPR')
+    # plt.ylabel('TPR')
+    # plt.savefig('output/ROC.pdf')
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+
+    # Required arguments
+    parser.add_argument('data_path', type=str,
+                        help='Path to the input dataset')
+    parser.add_argument('modeldir', type=str,
+                        help='Where to save the model')
+    parser.add_argument('submission_path', type=str,
+                        help='Where to save the model')
+
+    args = parser.parse_args()
+    main(args)

scripts/model.py

@@ -1,72 +1,26 @@
-import os
-import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
 
-import tensorflow as tf
-from tensorflow import keras
-from tensorflow.keras import layers
 
-from matplotlib import pyplot as plt
-from sklearn import metrics
-from sklearn.model_selection import train_test_split
+class ADClassifier(nn.Module):
 
+    def __init__(self):
+        super(ADClassifier, self).__init__()
 
-def transformer_encoder(inputs, head_size, num_heads, ff_dim, dropout=0):
-    """
-    ## Build the model
+        self.n_head = 1
+        self.flat1 = nn.Flatten()
+        self.encoder_layer = nn.TransformerEncoderLayer(d_model=400, nhead=self.n_head)
+        self.fc1 = nn.Linear(400*self.n_head, 400)
+        self.flat2 = nn.Flatten()
+        self.fc2 = nn.Linear(400, 1)
 
-    Our model processes a tensor of shape `(batch size, sequence length, features)`,
-    where `sequence length` is the number of time steps and `features` is each input
-    timeseries.
+    def forward(self, x):
 
-    We include residual connections, layer normalization, and dropout.
-    The resulting layer can be stacked multiple times.
+        x = self.flat1(x)
+        x = self.encoder_layer(x)
+        x = F.relu(self.fc1(x))
+        x = self.flat2(x)
+        x = F.sigmoid(self.fc2(x))
 
-    The projection layers are implemented through `keras.layers.Conv1D`.
-
-    """
-    # Attention and Normalization
-    x = layers.MultiHeadAttention(
-        key_dim=head_size, num_heads=num_heads, dropout=dropout
-    )(inputs, inputs)
-    x = layers.Dropout(dropout)(x)
-    x = layers.LayerNormalization(epsilon=1e-6)(x)
-    res = x + inputs
-
-    # Feed Forward Part
-    x = layers.Conv1D(filters=ff_dim, kernel_size=1, activation="relu")(res)
-    x = layers.Dropout(dropout)(x)
-    x = layers.Conv1D(filters=inputs.shape[-1], kernel_size=1)(x)
-    x = layers.LayerNormalization(epsilon=1e-6)(x)
-    return x + res
-
-
-def build_model(
-    input_shape,
-    head_size,
-    num_heads,
-    ff_dim,
-    num_transformer_blocks,
-    mlp_units,
-    dropout=0,
-    mlp_dropout=0,
-    ):
-    """
-    The main part of our model is now complete. We can stack multiple of those
-    `transformer_encoder` blocks and we can also proceed to add the final
-    Multi-Layer Perceptron classification head. Apart from a stack of `Dense`
-    layers, we need to reduce the output tensor of the `TransformerEncoder` part of
-    our model down to a vector of features for each data point in the current
-    batch. A common way to achieve this is to use a pooling layer. For
-    this example, a `GlobalAveragePooling1D` layer is sufficient.
-    """
-    inputs = keras.Input(shape=input_shape)
-    x = inputs
-    for _ in range(num_transformer_blocks):
-        x = transformer_encoder(x, head_size, num_heads, ff_dim, dropout)
-
-    x = layers.GlobalAveragePooling1D(data_format="channels_first")(x)
-    for dim in mlp_units:
-        x = layers.Dense(dim, activation="relu")(x)
-        x = layers.Dropout(mlp_dropout)(x)
-    outputs = layers.Dense(1, activation="sigmoid")(x)
-    return keras.Model(inputs, outputs)
+        return x