Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

1478 Fix TorchScript issue in AHnet #1479

Merged
merged 3 commits into from
Jan 21, 2021
Merged

1478 Fix TorchScript issue in AHnet #1479

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5e3d575
Merge pull request #13 from Project-MONAI/master
Nic-Ma Jan 21, 2021
7c9d4dc
[DLMED] fix TorchScript issue in AHNet
Nic-Ma Jan 21, 2021
8aa2acf
[DLMED] add test cases
Nic-Ma Jan 21, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion monai/networks/nets/ahnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -371,6 +371,7 @@ def __init__(
self.pool_type = pool_type
self.spatial_dims = spatial_dims
self.psp_block_num = psp_block_num
self.psp = None

if spatial_dims not in [2, 3]:
raise AssertionError("spatial_dims can only be 2 or 3.")
Expand Down Expand Up @@ -510,7 +511,7 @@ def forward(self, x):

sum4 = self.up3(d3) + conv_x
d4 = self.dense4(sum4)
if self.psp_block_num > 0:
if self.psp_block_num > 0 and self.psp is not None:
psp = self.psp(d4)
x = torch.cat((psp, d4), dim=1)
else:
Expand Down
5 changes: 5 additions & 0 deletions tests/test_ahnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,9 +191,14 @@ def test_ahnet_shape_3d(self, input_param, input_shape, expected_shape):

@skip_if_quick
def test_script(self):
# test 2D network
net = AHNet(spatial_dims=2, out_channels=2)
test_data = torch.randn(1, 1, 128, 64)
test_script_save(net, test_data)
# test 3D network
net = AHNet(spatial_dims=3, out_channels=2, psp_block_num=0, upsample_mode="nearest")
test_data = torch.randn(1, 1, 32, 32, 64)
test_script_save(net, test_data)


class TestAHNETWithPretrain(unittest.TestCase):
Expand Down