Skip to content

Commit 68ea7ce

Browse files
RMeliRMeli
authored
Merge pull request #16 from RMeli/highmodels
Add high-resolution models
2 parents 3c5da04 + d1fa671 commit 68ea7ce

2 files changed

Lines changed: 315 additions & 16 deletions

File tree

‎gnina/models.py‎

Lines changed: 280 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -588,7 +588,6 @@ def __init__(
588588
super().__init__()
589589

590590
self.input_dims = input_dims
591-
self.predict_affinity = affinity
592591

593592
features: OrderedDict[str, nn.Module] = OrderedDict(
594593
[
@@ -742,8 +741,6 @@ def forward(self, x):
742741
# Global max pooling reduced spatial dimensions to single value
743742
x = x.view(-1, self.features_out_size)
744743

745-
pose_raw = self.pose(x)
746-
747744
pose_raw = self.pose(x)
748745
pose_log = F.log_softmax(pose_raw, dim=1)
749746

@@ -787,19 +784,16 @@ def __init__(
787784
super().__init__(input_dims, num_blocks, num_block_features, num_block_convs)
788785

789786
# Linear layer for binding affinity prediction
790-
if self.predict_affinity:
791-
self.affinity = nn.Sequential(
792-
OrderedDict(
793-
[
794-
(
795-
"affinity_output",
796-
nn.Linear(
797-
in_features=self.features_out_size, out_features=1
798-
),
799-
)
800-
]
801-
)
787+
self.affinity = nn.Sequential(
788+
OrderedDict(
789+
[
790+
(
791+
"affinity_output",
792+
nn.Linear(in_features=self.features_out_size, out_features=1),
793+
)
794+
]
802795
)
796+
)
803797

804798
# Xavier initialization for convolutional and linear layers
805799
for m in self.modules():
@@ -824,6 +818,275 @@ def forward(self, x):
824818
x = x.view(-1, self.features_out_size)
825819

826820
pose_raw = self.pose(x)
821+
pose_log = F.log_softmax(pose_raw, dim=1)
822+
823+
affinity = self.affinity(x)
824+
# Squeeze last (dummy) dimension of affinity prediction
825+
# This allows to match the shape (batch_size,) of the target tensor
826+
return pose_log, affinity.squeeze(-1)
827+
828+
829+
class HiResPose(nn.Module):
830+
"""
831+
GNINA HiResPose model architecture.
832+
833+
Parameters
834+
----------
835+
input_dims: tuple
836+
Model input dimensions (channels, depth, height, width)
837+
838+
Notes
839+
-----
840+
This architecture was translated from the following Caffe model:
841+
842+
https://github.com/gnina/models/blob/master/crossdocked_paper/hires_pose.model
843+
844+
The main difference is that the PyTorch implementation resurns the log softmax.
845+
846+
This model is implemented only for multi-task pose and affinity prediction.
847+
"""
848+
849+
def __init__(self, input_dims: Tuple):
850+
851+
super().__init__()
852+
853+
self.input_dims = input_dims
854+
855+
self.features = nn.Sequential(
856+
OrderedDict(
857+
[
858+
# unit1
859+
(
860+
"unit1_conv",
861+
nn.Conv3d(
862+
in_channels=input_dims[0],
863+
out_channels=32,
864+
kernel_size=3,
865+
stride=1,
866+
padding=1,
867+
),
868+
),
869+
("unit1_func", nn.ReLU()),
870+
# unit2
871+
("unit2_pool", nn.MaxPool3d(kernel_size=2, stride=2)),
872+
(
873+
"unit2_conv",
874+
nn.Conv3d(
875+
in_channels=32,
876+
out_channels=64,
877+
kernel_size=3,
878+
stride=1,
879+
padding=1,
880+
),
881+
),
882+
("unit2_func", nn.ReLU()),
883+
# unit3
884+
("unit3_pool", nn.MaxPool3d(kernel_size=2, stride=2)),
885+
(
886+
"unit3_conv",
887+
nn.Conv3d(
888+
in_channels=64,
889+
out_channels=128,
890+
kernel_size=3,
891+
stride=1,
892+
padding=1,
893+
),
894+
),
895+
("unit3_func", nn.ReLU()),
896+
]
897+
)
898+
)
899+
900+
# Two MaxPool3d layers with kernel_size=2 and stride=2
901+
# Spatial dimensions are halved at each pooling step
902+
self.features_out_size = (
903+
input_dims[1] // 4 * input_dims[2] // 4 * input_dims[3] // 4 * 128
904+
)
905+
906+
# Linear layer for pose prediction
907+
self.pose = nn.Sequential(
908+
OrderedDict(
909+
[
910+
(
911+
"pose_output",
912+
nn.Linear(in_features=self.features_out_size, out_features=2),
913+
)
914+
]
915+
)
916+
)
917+
918+
# Linear layer for binding affinity prediction
919+
self.affinity = nn.Sequential(
920+
OrderedDict(
921+
[
922+
(
923+
"affinity_output",
924+
nn.Linear(in_features=self.features_out_size, out_features=1),
925+
)
926+
]
927+
)
928+
)
929+
930+
# Xavier initialization for convolutional and linear layers
931+
for m in self.modules():
932+
if isinstance(m, nn.Conv3d) or isinstance(m, nn.Linear):
933+
nn.init.xavier_uniform_(m.weight.data)
934+
935+
def forward(self, x: torch.Tensor):
936+
"""
937+
Parameters
938+
----------
939+
x: torch.Tensor
940+
Input tensor
941+
942+
Notes
943+
-----
944+
The pose score is the log softmax of the output of the last linear layer.
945+
"""
946+
x = self.features(x)
947+
948+
print("FEATURES SHAPE:", x.shape)
949+
950+
# Reshape based on number of channels
951+
# Global max pooling reduced spatial dimensions to single value
952+
x = x.view(-1, self.features_out_size)
953+
954+
pose_raw = self.pose(x)
955+
pose_log = F.log_softmax(pose_raw, dim=1)
956+
957+
affinity = self.affinity(x)
958+
# Squeeze last (dummy) dimension of affinity prediction
959+
# This allows to match the shape (batch_size,) of the target tensor
960+
return pose_log, affinity.squeeze(-1)
961+
962+
963+
class HiResAffinity(nn.Module):
964+
"""
965+
GNINA HiResAffinity model architecture.
966+
967+
Parameters
968+
----------
969+
input_dims: tuple
970+
Model input dimensions (channels, depth, height, width)
971+
972+
Notes
973+
-----
974+
This architecture was translated from the following Caffe model:
975+
976+
https://github.com/gnina/models/blob/master/crossdocked_paper/hires_pose.model
977+
978+
The main difference is that the PyTorch implementation resurns the log softmax.
979+
980+
This model is implemented only for multi-task pose and affinity prediction.
981+
"""
982+
983+
def __init__(self, input_dims: Tuple):
984+
985+
super().__init__()
986+
987+
self.input_dims = input_dims
988+
989+
self.features = nn.Sequential(
990+
OrderedDict(
991+
[
992+
# unit1
993+
(
994+
"unit1_conv",
995+
nn.Conv3d(
996+
in_channels=input_dims[0],
997+
out_channels=32,
998+
kernel_size=3,
999+
stride=1,
1000+
padding=1,
1001+
),
1002+
),
1003+
("unit1_func", nn.ReLU()),
1004+
# unit2
1005+
(
1006+
"unit2_conv",
1007+
nn.Conv3d(
1008+
in_channels=32,
1009+
out_channels=64,
1010+
kernel_size=3,
1011+
stride=1,
1012+
padding=1,
1013+
),
1014+
),
1015+
("unit2_func", nn.ReLU()),
1016+
# unit3
1017+
("unit3_pool", nn.AvgPool3d(kernel_size=8, stride=8)),
1018+
(
1019+
"unit3_conv",
1020+
nn.Conv3d(
1021+
in_channels=64,
1022+
out_channels=128,
1023+
kernel_size=5,
1024+
stride=1,
1025+
padding=2,
1026+
),
1027+
),
1028+
("unit3_func", nn.ELU(alpha=1.0)),
1029+
# unit5 (following original naming convention)
1030+
("unit5_pool", nn.MaxPool3d(kernel_size=4, stride=4)),
1031+
]
1032+
)
1033+
)
1034+
1035+
self.features_out_size = (
1036+
input_dims[1]
1037+
// (8 * 4)
1038+
* input_dims[2]
1039+
// (8 * 4)
1040+
* input_dims[3]
1041+
// (8 * 4)
1042+
* 128
1043+
)
1044+
1045+
# Linear layer for pose prediction
1046+
self.pose = nn.Sequential(
1047+
OrderedDict(
1048+
[
1049+
(
1050+
"pose_output",
1051+
nn.Linear(in_features=self.features_out_size, out_features=2),
1052+
)
1053+
]
1054+
)
1055+
)
1056+
1057+
# Linear layer for binding affinity prediction
1058+
self.affinity = nn.Sequential(
1059+
OrderedDict(
1060+
[
1061+
(
1062+
"affinity_output",
1063+
nn.Linear(in_features=self.features_out_size, out_features=1),
1064+
)
1065+
]
1066+
)
1067+
)
1068+
1069+
# Xavier initialization for convolutional and linear layers
1070+
for m in self.modules():
1071+
if isinstance(m, nn.Conv3d) or isinstance(m, nn.Linear):
1072+
nn.init.xavier_uniform_(m.weight.data)
1073+
1074+
def forward(self, x: torch.Tensor):
1075+
"""
1076+
Parameters
1077+
----------
1078+
x: torch.Tensor
1079+
Input tensor
1080+
1081+
Notes
1082+
-----
1083+
The pose score is the log softmax of the output of the last linear layer.
1084+
"""
1085+
x = self.features(x)
1086+
1087+
# Reshape based on number of channels
1088+
# Global max pooling reduced spatial dimensions to single value
1089+
x = x.view(-1, self.features_out_size)
8271090

8281091
pose_raw = self.pose(x)
8291092
pose_log = F.log_softmax(pose_raw, dim=1)
@@ -841,4 +1104,6 @@ def forward(self, x):
8411104
("default2018", True): Default2018Affinity,
8421105
("dense", False): DensePose,
8431106
("dense", True): DenseAffinity,
1107+
("hires_pose", True): HiResPose,
1108+
("hires_affinity", True): HiResAffinity,
8441109
}

‎tests/test_models.py‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,27 @@ def dims():
1414
return (12, 24, 24, 24)
1515

1616

17+
@pytest.fixture
18+
def dims_big():
19+
"""
20+
Less channels, but bigger spatial dimensions (like original input)
21+
"""
22+
return (3, 48, 48, 48)
23+
24+
1725
@pytest.fixture
1826
def x(batch_size, dims, device):
1927
return torch.normal(mean=0, std=1, size=(batch_size, *dims), device=device)
2028

2129

30+
@pytest.fixture
31+
def x_big(batch_size, dims_big, device):
32+
"""
33+
HiResAffinity has as an aggressive pooling, so we need the standard input size.
34+
"""
35+
return torch.normal(mean=0, std=1, size=(batch_size, *dims_big), device=device)
36+
37+
2238
@pytest.mark.parametrize("model", ["default2017", "default2018", "dense"])
2339
def test_forward_pose(batch_size, dims, x, device, model):
2440
"""
@@ -30,10 +46,14 @@ def test_forward_pose(batch_size, dims, x, device, model):
3046
assert pose_raw.shape == (batch_size, 2)
3147

3248

33-
@pytest.mark.parametrize("model", ["default2017", "default2018", "dense"])
49+
@pytest.mark.parametrize("model", ["default2017", "default2018", "dense", "hires_pose"])
3450
def test_forward_affinity(batch_size, dims, x, device, model):
3551
"""
3652
Test forward pass of models for pose and binding affinity prediction.
53+
54+
Notes
55+
-----
56+
All models but :code:`hires_affinity`, which requires larger spatial dimensions.
3757
"""
3858
m = models_dict[(model, True)](input_dims=dims).to(device)
3959
pose_log, affinity = m(x)
@@ -42,6 +62,20 @@ def test_forward_affinity(batch_size, dims, x, device, model):
4262
assert affinity.shape == (batch_size,)
4363

4464

65+
@pytest.mark.parametrize(
66+
"model", ["default2017", "default2018", "dense", "hires_pose", "hires_affinity"]
67+
)
68+
def test_forward_affinity_big(batch_size, dims_big, x_big, device, model):
69+
"""
70+
Test forward pass of models for pose and binding affinity prediction.
71+
"""
72+
m = models_dict[(model, True)](input_dims=dims_big).to(device)
73+
pose_log, affinity = m(x_big)
74+
75+
assert pose_log.shape == (batch_size, 2)
76+
assert affinity.shape == (batch_size,)
77+
78+
4579
@pytest.mark.parametrize("num_block_convs", [1, 4])
4680
@pytest.mark.parametrize("num_block_features", [2, 16])
4781
def test_denseblock_forward_small(

0 commit comments

Comments
 (0)