ChemicalX

Author: benedekrozemberczki

chemicalx/models/deepsynergy.py

@@ -38,7 +38,7 @@ def forward(
         drug_features_right: torch.FloatTensor,
     ) -> torch.FloatTensor:
 
-        hidden = torch.cat([drug_features_left, drug_features_right, context_features], dim=1)
+        hidden = torch.cat([context_features, drug_features_left, drug_features_right], dim=1)
         hidden = self.encoder(hidden)
         hidden = F.relu(hidden)
         hidden = self.hidden_first(hidden)

tests/unit/test_models.py

@@ -78,22 +78,24 @@ def test_DeepDrug(self):
     def test_DeepSynergy(self):
 
         model = DeepSynergy(
-            context_channels=64,
-            drug_channels=32,
+            context_channels=112,
+            drug_channels=256,
             input_hidden_channels=32,
             middle_hidden_channels=16,
             final_hidden_channels=16,
             dropout_rate=0.5,
         )
 
         optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=0.0001)
+        model.train()
         loss = torch.nn.BCELoss()
         for batch in self.generator:
             optimizer.zero_grad()
             prediction = model(batch.context_features, batch.drug_features_left, batch.drug_features_right)
             output = loss(prediction, batch.labels)
             output.backward()
             optimizer.step()
+            assert prediction.shape[0] == batch.labels.shape[0]
 
     def test_DeepDDS(self):
         model = DeepDDS(x=2)