Add TorchScript compatible test for all the networks

[Nic-Ma]

Is your feature request related to a problem? Please describe.
In order to make sure all the networks can be exported and loaded in TorchScript, need to add more unit tests for them.

[Nic-Ma]

Hi @ericspod ,

I am working on a PR #1121 to address this ticket, but got an TorchScript error in AHNet:
https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/nets/ahnet.py#L302

======================================================================
ERROR: test_ahnet_script (__main__.TestAHNET)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "tests/test_ahnet.py", line 145, in test_ahnet_script
    out_orig, out_reloaded = test_script_save(net, test_data)
  File "/workspace/data/medical/MONAI/tests/utils.py", line 100, in test_script_save
    scripted = torch.jit.script(net)
  File "/opt/conda/lib/python3.6/site-packages/torch/jit/__init__.py", line 1516, in script
    return torch.jit._recursive.create_script_module(obj, torch.jit._recursive.infer_methods_to_compile)
  File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 318, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 372, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/opt/conda/lib/python3.6/site-packages/torch/jit/__init__.py", line 1900, in _construct
    init_fn(script_module)
  File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 353, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, infer_methods_to_compile)
  File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 376, in create_script_module_impl
    create_methods_from_stubs(concrete_type, stubs)
  File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 292, in create_methods_from_stubs
    concrete_type._create_methods(defs, rcbs, defaults)
RuntimeError:
cannot statically infer the expected size of a list in this context:
  File "/workspace/data/medical/MONAI/monai/networks/nets/ahnet.py", line 302
            x8 = self.up8(self.proj8(self.pool8(x)))
        else:
            interpolate_size = tuple(x.size()[2:])
                               ~~~~~~~~~~~~~~~~ <--- HERE
            x64 = F.interpolate(
                self.proj64(self.pool64(x)),

Because this line is not static shape for TorchScript, what do you think we can do for this kind case?

Thanks.

[Nic-Ma]

Hi @ericspod ,

Seems several our networks contain dynamic parameters, so torch.jit.script can't work, can we use torch.jit.trace to test it?
Thanks.

[ericspod]

One thing to do is replace interpolate_size = tuple(x.size()[2:]) with interpolate_size = x.shape[2:] and this should work. In the PSP class the members u8, u16, u32, u64 are not always created, Torchscript won't like this so if upsample mode isn't transpose these must be set to nn.Identity so that they are always there. The line self.dropout_prob = 0 should also be self.dropout_prob = 0.0 for type reasons.

[Nic-Ma]

Cool, thanks for your help!