Merge pull request #144 from kathyxchen/heartenn-updates

Add HeartENN arch to models dir

Author: kathyxchen

models/README.md

@@ -1,5 +1,5 @@
 # Models directory
 
-We will provide PyTorch modules of different model architectures here.
+We will provide example PyTorch modules of different model architectures here.
 
 If you have any particular architecture requests, please let us know. Thanks!

models/heartenn.py

@@ -0,0 +1,89 @@
+"""
+HeartENN architecture (Richter et al., 2020).
+"""
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class HeartENN(nn.Module):
+    def __init__(self, sequence_length, n_genomic_features):
+        """
+        Parameters
+        ----------
+        sequence_length : int
+            Length of sequence context on which to train.
+        n_genomic_features : int
+            The number of chromatin features to predict.
+
+        Attributes
+        ----------
+        conv_net : torch.nn.Sequential
+        classifier : torch.nn.Sequential
+
+        """
+        super(HeartENN, self).__init__()
+        conv_kernel_size = 8
+        pool_kernel_size = 4
+
+        self.conv_net = nn.Sequential(
+            nn.Conv1d(4, 60, kernel_size=conv_kernel_size),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(60, 60, kernel_size=conv_kernel_size),
+            nn.ReLU(inplace=True),
+            nn.MaxPool1d(
+                kernel_size=pool_kernel_size, stride=pool_kernel_size),
+            nn.BatchNorm1d(60),
+
+            nn.Conv1d(60, 80, kernel_size=conv_kernel_size),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(80, 80, kernel_size=conv_kernel_size),
+            nn.ReLU(inplace=True),
+            nn.MaxPool1d(
+                kernel_size=pool_kernel_size, stride=pool_kernel_size),
+            nn.BatchNorm1d(80),
+            nn.Dropout(p=0.4),
+
+            nn.Conv1d(80, 240, kernel_size=conv_kernel_size),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(240, 240, kernel_size=conv_kernel_size),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm1d(240),
+            nn.Dropout(p=0.6))
+
+        reduce_by = 2 * (conv_kernel_size - 1)
+        pool_kernel_size = float(pool_kernel_size)
+        self._n_channels = int(
+            np.floor(
+                (np.floor(
+                    (sequence_length - reduce_by) / pool_kernel_size)
+                 - reduce_by) / pool_kernel_size)
+            - reduce_by)
+        self.classifier = nn.Sequential(
+            nn.Linear(240 * self._n_channels, n_genomic_features),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm1d(n_genomic_features),
+            nn.Linear(n_genomic_features, n_genomic_features),
+            nn.Sigmoid())
+
+    def forward(self, x):
+        """Forward propagation of a batch.i
+
+        """
+        for layer in self.conv_net.children():
+            if isinstance(layer, nn.Conv1d):
+                layer.weight.data.renorm_(2, 0, 0.9)
+        for layer in self.classifier.children():
+            if isinstance(layer, nn.Linear):
+                layer.weight.data.renorm_(2, 0, 0.9)
+        out = self.conv_net(x)
+        reshape_out = out.view(out.size(0), 240 * self._n_channels)
+        predict = self.classifier(reshape_out)
+        return predict
+
+def criterion():
+    return nn.BCELoss()
+
+def get_optimizer(lr):
+    return (torch.optim.SGD,
+            {"lr": lr, "weight_decay": 1e-6, "momentum": 0.9})

selene_sdk/utils/utils.py

@@ -7,7 +7,6 @@
 from collections import OrderedDict
 import logging
 import sys
-import traceback
 
 import numpy as np
 
@@ -99,32 +98,30 @@ def load_model_from_state_dict(state_dict, model):
     model_keys = model.state_dict().keys()
     state_dict_keys = state_dict.keys()
 
-    new_state_dict = OrderedDict()
-
     if len(model_keys) != len(state_dict_keys):
-        raise ValueError("State dict does not have the same "
-            "number of modules as the specified model "
-            "architecture. Please check whether you are using "
-            "the expected model architecture and that your PyTorch "
-            "version matches the version in which the loaded model "
-            "was trained.\n\n"
-            "\tExpected modules:\n\t{0}\n\n"
-            "\tModules in the loaded model weights:\n\t{1}\n".format(
-                model_keys, state_dict_keys))
+        try:
+            model.load_state_dict(state_dict, strict=False)
+            return model
+        except Exception as e:
+            raise ValueError("Loaded state dict does not match the model "
+                "architecture specified - please check that you are "
+                "using the correct architecture file and parameters.\n\n"
+                "{0}".format(e))
 
+    new_state_dict = OrderedDict()
     for (k1, k2) in zip(model_keys, state_dict_keys):
         value = state_dict[k2]
         try:
             new_state_dict[k1] = value
-        except Exception:
+        except Exception as e:
             raise ValueError(
                 "Failed to load weight from module {0} in model weights "
-                "into model architecture module {1}. (If module name "
+                "into model architecture module {1}. (If module name has "
                 "an additional prefix `model.` it is because the model is "
                 "wrapped in `selene_sdk.utils.NonStrandSpecific`. This "
                 "error was raised because the underlying module does "
                 "not match that expected by the loaded model:\n"
-                "{2}".format(k2, k1, traceback.print_exc()))
+                "{2}".format(k2, k1, e))
     model.load_state_dict(new_state_dict)
     return model
 

selene_sdk/version.py

@@ -1 +1 @@
-__version__ = "0.4.7"
+__version__ = "0.4.8"

setup.py

@@ -25,7 +25,7 @@
 cmdclass = {'build_ext': build_ext}
 
 setup(name="selene-sdk",
-      version="0.4.7",
+      version="0.4.8",
       long_description=long_description,
       long_description_content_type='text/markdown',
       description=("framework for developing sequence-level "