Repository: pytorch
Issue: multi-input and multi-output support for nn.Module
Labels: module: nn,triaged

## 🚀 Feature
after a model is constructed, whichever module(layer) you are at right now, you can get its input layer list and output layer list.

## Motivation

I'm writing a demo to go through the whole network and print its structure. The problem is that I cannot get the connection info(ie. when you are at one spesifical child module, there is no way that you can get all its parent modules or subchild modules), especially for multi-input and multi-output module as well as multi branch module like ResNet.

## Alternatives

there is no perfect solution right now.
I've tried the following methods:
- using Module's hook machinism (you can only get the input and output tensors, but not input layers and output layers)
- writing my own show_summary() function and it will work only when several condition met.

## Additional context
here is my own struggle to print the structure of ResNet18 (for image classification on cifar10):
the result is:
```
         model.conv1(Conv2d) (128, 3, 32, 32)->[128, 64, 32, 32]
	 model.bn1(BatchNorm2d) [128, 64, 32, 32]->[128, 64, 32, 32]
	 model.relu(ReLU) [128, 64, 32, 32]->[128, 64, 32, 32]
	 model.group1(Sequential) [128, 64, 32, 32]->[128, 64, 32, 32]
		 model.group1.0(BaseBlock) [128, 64, 32, 32]->[128, 64, 32, 32]
			 model.group1.0.conv1(Conv2d) [128, 64, 32, 32]->[128, 64, 32, 32]
			 model.group1.0.bn1(BatchNorm2d) [128, 64, 32, 32]->[128, 64, 32, 32]
			 model.group1.0.relu(ReLU) [128, 64, 32, 32]->[128, 64, 32, 32]
			 model.group1.0.conv2(Conv2d) [128, 64, 32, 32]->[128, 64, 32, 32]
			 model.group1.0.bn2(BatchNorm2d) [128, 64, 32, 32]->[128, 64, 32, 32]
		 model.group1.1(BaseBlock) [128, 64, 32, 32]->[128, 64, 32, 32]
			 model.group1.1.conv1(Conv2d) [128, 64, 32, 32]->[128, 64, 32, 32]
			 model.group1.1.bn1(BatchNorm2d) [128, 64, 32, 32]->[128, 64, 32, 32]
			 model.group1.1.relu(ReLU) [128, 64, 32, 32]->[128, 64, 32, 32]
			 model.group1.1.conv2(Conv2d) [128, 64, 32, 32]->[128, 64, 32, 32]
			 model.group1.1.bn2(BatchNorm2d) [128, 64, 32, 32]->[128, 64, 32, 32]
	 model.group2(Sequential) [128, 64, 32, 32]->[128, 128, 16, 16]
		 model.group2.0(BaseBlock) [128, 64, 32, 32]->[128, 128, 16, 16]
			 model.group2.0.conv1(Conv2d) [128, 64, 32, 32]->[128, 128, 16, 16]
			 model.group2.0.bn1(BatchNorm2d) [128, 128, 16, 16]->[128, 128, 16, 16]
			 model.group2.0.relu(ReLU) [128, 128, 16, 16]->[128, 128, 16, 16]
			 model.group2.0.conv2(Conv2d) [128, 128, 16, 16]->[128, 128, 16, 16]
			 model.group2.0.bn2(BatchNorm2d) [128, 128, 16, 16]->[128, 128, 16, 16]
			 model.group2.0.shortcut(Sequential) [128, 64, 32, 32]->[128, 128, 16, 16]
				 model.group2.0.shortcut.0(Conv2d) [128, 64, 32, 32]->[128, 128, 16, 16]
				 model.group2.0.shortcut.1(BatchNorm2d) [128, 128, 16, 16]->[128, 128, 16, 16]
		 model.group2.1(BaseBlock) [128, 128, 16, 16]->[128, 128, 16, 16]
			 model.group2.1.conv1(Conv2d) [128, 128, 16, 16]->[128, 128, 16, 16]
			 model.group2.1.bn1(BatchNorm2d) [128, 128, 16, 16]->[128, 128, 16, 16]
			 model.group2.1.relu(ReLU) [128, 128, 16, 16]->[128, 128, 16, 16]
			 model.group2.1.conv2(Conv2d) [128, 128, 16, 16]->[128, 128, 16, 16]
			 model.group2.1.bn2(BatchNorm2d) [128, 128, 16, 16]->[128, 128, 16, 16]
	 model.group3(Sequential) [128, 128, 16, 16]->[128, 256, 8, 8]
		 model.group3.0(BaseBlock) [128, 128, 16, 16]->[128, 256, 8, 8]
			 model.group3.0.conv1(Conv2d) [128, 128, 16, 16]->[128, 256, 8, 8]
			 model.group3.0.bn1(BatchNorm2d) [128, 256, 8, 8]->[128, 256, 8, 8]
			 model.group3.0.relu(ReLU) [128, 256, 8, 8]->[128, 256, 8, 8]
			 model.group3.0.conv2(Conv2d) [128, 256, 8, 8]->[128, 256, 8, 8]
			 model.group3.0.bn2(BatchNorm2d) [128, 256, 8, 8]->[128, 256, 8, 8]
			 model.group3.0.shortcut(Sequential) [128, 128, 16, 16]->[128, 256, 8, 8]
				 model.group3.0.shortcut.0(Conv2d) [128, 128, 16, 16]->[128, 256, 8, 8]
				 model.group3.0.shortcut.1(BatchNorm2d) [128, 256, 8, 8]->[128, 256, 8, 8]
		 model.group3.1(BaseBlock) [128, 256, 8, 8]->[128, 256, 8, 8]
			 model.group3.1.conv1(Conv2d) [128, 256, 8, 8]->[128, 256, 8, 8]
			 model.group3.1.bn1(BatchNorm2d) [128, 256, 8, 8]->[128, 256, 8, 8]
			 model.group3.1.relu(ReLU) [128, 256, 8, 8]->[128, 256, 8, 8]
			 model.group3.1.conv2(Conv2d) [128, 256, 8, 8]->[128, 256, 8, 8]
			 model.group3.1.bn2(BatchNorm2d) [128, 256, 8, 8]->[128, 256, 8, 8]
	 model.group4(Sequential) [128, 256, 8, 8]->[128, 512, 4, 4]
		 model.group4.0(BaseBlock) [128, 256, 8, 8]->[128, 512, 4, 4]
			 model.group4.0.conv1(Conv2d) [128, 256, 8, 8]->[128, 512, 4, 4]
			 model.group4.0.bn1(BatchNorm2d) [128, 512, 4, 4]->[128, 512, 4, 4]
			 model.group4.0.relu(ReLU) [128, 512, 4, 4]->[128, 512, 4, 4]
			 model.group4.0.conv2(Conv2d) [128, 512, 4, 4]->[128, 512, 4, 4]
			 model.group4.0.bn2(BatchNorm2d) [128, 512, 4, 4]->[128, 512, 4, 4]
			 model.group4.0.shortcut(Sequential) [128, 256, 8, 8]->[128, 512, 4, 4]
				 model.group4.0.shortcut.0(Conv2d) [128, 256, 8, 8]->[128, 512, 4, 4]
				 model.group4.0.shortcut.1(BatchNorm2d) [128, 512, 4, 4]->[128, 512, 4, 4]
		 model.group4.1(BaseBlock) [128, 512, 4, 4]->[128, 512, 4, 4]
			 model.group4.1.conv1(Conv2d) [128, 512, 4, 4]->[128, 512, 4, 4]
			 model.group4.1.bn1(BatchNorm2d) [128, 512, 4, 4]->[128, 512, 4, 4]
			 model.group4.1.relu(ReLU) [128, 512, 4, 4]->[128, 512, 4, 4]
			 model.group4.1.conv2(Conv2d) [128, 512, 4, 4]->[128, 512, 4, 4]
			 model.group4.1.bn2(BatchNorm2d) [128, 512, 4, 4]->[128, 512, 4, 4]
	 model.avg_pool(AvgPool2d) [128, 512, 4, 4]->[128, 512, 1, 1]
	 model.flatten(Flatten) [128, 512, 1, 1]->[128, 512]
	 model.linear(Linear) [128, 512]->[128, 10]
```

the model is defined like this:
```python
import torch
import torch.nn as nn


BATCH_NORMALIZATION_MOMENTUM = 0.1


class BaseBlock(nn.Module):
    expansion = 1

    def __init__(self, input_channel, output_channel, stride=1, shortcut=None):
        super(BaseBlock, self).__init__()
        # 定义BaseBlock中要用到的层
        self.conv1 = nn.Conv2d(input_channel, output_channel, kernel_size=(3, 3), stride=(stride, stride), padding=(1, 1), bias=False)
        self.bn1 = nn.BatchNorm2d(output_channel, momentum=BATCH_NORMALIZATION_MOMENTUM)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(output_channel, output_channel, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        self.bn2 = nn.BatchNorm2d(output_channel, momentum=BATCH_NORMALIZATION_MOMENTUM)
        # 储存传入BaseBlock的参数
        self.shortcut = shortcut
        self.stride = stride

    def forward(self, input):
        output = self.conv1(input)
        output = self.bn1(output)
        output = self.relu(output)
        output = self.conv2(output)
        output = self.bn2(output)
        if self.shortcut is not None:
            skip = self.shortcut(input)
        else:
            skip = input
        output += skip
        output = self.relu(output)
        return output


class ResNet18(nn.Module):
    def __init__(self):
        super(ResNet18, self).__init__()
        # conv1
        self.conv1 = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        self.bn1 = nn.BatchNorm2d(64, momentum=BATCH_NORMALIZATION_MOMENTUM)
        self.relu = nn.ReLU()

        # group1
        block1_1_shortcut = None
        block1_1 = BaseBlock(input_channel=64, output_channel=64, stride=1, shortcut=block1_1_shortcut)
        block1_2 = BaseBlock(input_channel=64, output_channel=64, stride=1, shortcut=None)
        self.group1 = nn.Sequential(block1_1, block1_2)

        # group2
        block2_1_shortcut = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), padding=(0, 0), bias=False),
            nn.BatchNorm2d(128, momentum=BATCH_NORMALIZATION_MOMENTUM),
        )
        block2_1 = BaseBlock(input_channel=64, output_channel=128, stride=2, shortcut=block2_1_shortcut)
        block2_2 = BaseBlock(input_channel=128, output_channel=128, stride=1, shortcut=None)
        self.group2 = nn.Sequential(block2_1, block2_2)

        # group3
        block3_1_shortcut = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), padding=(0, 0), bias=False),
            nn.BatchNorm2d(256, momentum=BATCH_NORMALIZATION_MOMENTUM),
        )
        block3_1 = BaseBlock(input_channel=128, output_channel=256, stride=2, shortcut=block3_1_shortcut)
        block3_2 = BaseBlock(input_channel=256, output_channel=256, stride=1, shortcut=None)
        self.group3 = nn.Sequential(block3_1, block3_2)

        # group4
        block4_1_shortcut = nn.Sequential(
            nn.Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), padding=(0, 0), bias=False),
            nn.BatchNorm2d(512, momentum=BATCH_NORMALIZATION_MOMENTUM),
        )
        block4_1 = BaseBlock(input_channel=256, output_channel=512, stride=2, shortcut=block4_1_shortcut)
        block4_2 = BaseBlock(input_channel=512, output_channel=512, stride=1, shortcut=None)
        self.group4 = nn.Sequential(block4_1, block4_2)

        self.avg_pool = nn.AvgPool2d(4)
        self.flatten = nn.Flatten()
        self.linear = nn.Linear(512, 10)

    def forward(self, input):
        output = self.conv1(input)
        output = self.bn1(output)
        output = self.relu(output)

        output = self.group1(output)
        output = self.group2(output)
        output = self.group3(output)
        output = self.group4(output)

        output = self.avg_pool(output)
        # output = output.view(output.size(0), -1)
        output = self.flatten(output)  # 使用flatten层来代替上一行的reshape操作
        output = self.linear(output)
        return output
```

and the show_summary() function is like:
```python
# show_summary()方法用于打印模型各层的输出，is_shortcut()方法用于判断某个module是否为分支模块
# 注意：此工具方法的使用限制如下：
#       - 单个Module仅支持这样的多分支的结构：
#       1 各个分支的输入相同
#       2 所有分支是并列的地位，即各分支的输入相同，各分支的输出在同一时刻合并，各分支在合并以后不允许后面继续拼接主分支的结构（即各分支在合并以后直接作为Module的输出）
#       3 如果模型结构与2冲突，应该将模型通过嵌套Module来实现复杂的分支结构
#       4 分支的命名中必须包含shortcut（该关键字可以通过is_shortcut方法的keyword参数来修改），例如并列分支的命名可以是shortcut1，shortcut2，shortcut3
def is_shortcut(name, keyword="shortcut"):
    return len(name.split('.')) > 0 and name.split('.')[-1] and keyword in name.split('.')[-1]


def show_summary(model, input_shape, prefix="model"):
    input_shape_dict = dict()
    current_shape_dict = dict()
    input_shape_dict[prefix] = input_shape
    current_shape_dict[prefix] = input_shape
    for name, module in model.named_modules():
        if isinstance(model, nn.Module) and not module == model:
            current_name = prefix + "." + name
            parent_name = current_name[0:current_name.rfind(".")]
            level = len(current_name.split(".")) - 1
            module_type = str(module.__class__).split(".")[-1].split("'")[0]
            # print("{} {}->{}".format(level, parent_name, current_name))
            temp_input_shape = input_shape_dict[parent_name] if is_shortcut(current_name) else current_shape_dict[parent_name]
            tensor_temp = torch.randn(*temp_input_shape)
            output_temp = module(tensor_temp)
            output_shape = [len(output_temp)] + list(output_temp[0].shape)
            print("{} {}({}) {}->{}".format("\t" * (level), current_name, module_type, temp_input_shape, output_shape))
            input_shape_dict[current_name] = temp_input_shape
            current_shape_dict[current_name] = temp_input_shape
            current_shape_dict[parent_name] = output_shape
```

and my function also works for other model like MobileNet. my demo link is [click here](https://github.com/duwangthefirst/playground/tree/master/models).


cc @albanD @mruberry @jbschlosser