docstr: 构建对应stage，即用来融合不同尺度的实现
  :param input_branches: 输入的分支数，每个分支对应一种尺度
  :param output_branches: 输出的分支数
  :param c: 输入的第一个分支通道数
docstr_tok: ["\u6784\u5efa\u5bf9\u5e94stage\uff0c\u5373\u7528\u6765\u878d\u5408\u4e0d\u540c\u5c3a\u5ea6\u7684\u5b9e\u73b0",":","param","input_branches",":","\u8f93\u5165\u7684\u5206\u652f\u6570\uff0c\u6bcf\u4e2a\u5206\u652f\u5bf9\u5e94\u4e00\u79cd\u5c3a\u5ea6",":","param","output_branches",":","\u8f93\u51fa\u7684\u5206\u652f\u6570",":","param","c",":","\u8f93\u5165\u7684\u7b2c\u4e00\u4e2a\u5206\u652f\u901a\u9053\u6570"]
code: def __init__(self, input_branches, output_branches, c):
      """
      构建对应stage，即用来融合不同尺度的实现
      :param input_branches: 输入的分支数，每个分支对应一种尺度
      :param output_branches: 输出的分支数
      :param c: 输入的第一个分支通道数
      """
      super().__init__()
      self.input_branches = input_branches
      self.output_branches = output_branches
  
      self.branches = nn.ModuleList()
      for i in range(self.input_branches):  # 每个分支上都先通过4个BasicBlock
          w = c * (2**i)  # 对应第i个分支的通道数
          branch = nn.Sequential(
              BasicBlock(w, w),
              BasicBlock(w, w),
              BasicBlock(w, w),
              BasicBlock(w, w),
          )
          self.branches.append(branch)
  
      self.fuse_layers = nn.ModuleList()  # 用于融合每个分支上的输出
      for i in range(self.output_branches):
          self.fuse_layers.append(nn.ModuleList())
          for j in range(self.input_branches):
              if i == j:
                  # 当输入、输出为同一个分支时不做任何处理
                  self.fuse_layers[-1].append(nn.Identity())
              elif i < j:
                  # 当输入分支j大于输出分支i时(即输入分支下采样率大于输出分支下采样率)，
                  # 此时需要对输入分支j进行通道调整以及上采样，方便后续相加
                  self.fuse_layers[-1].append(
                      nn.Sequential(
                          nn.Conv2d(
                              c * (2**j),
                              c * (2**i),
                              kernel_size=1,
                              stride=1,
                              bias=False,
                          ),
                          nn.BatchNorm2d(
                              c * (2**i), momentum=BN_MOMENTUM
                          ),
                          nn.Upsample(
                              scale_factor=2.0 ** (j - i),
                              mode="nearest",
                          ),
                      )
                  )
              else:  # i > j
                  # 当输入分支j小于输出分支i时(即输入分支下采样率小于输出分支下采样率)，
                  # 此时需要对输入分支j进行通道调整以及下采样，方便后续相加
                  # 注意，这里每次下采样2x都是通过一个3x3卷积层实现的，4x就是两个，8x就是三个，总共i-j个
                  ops = []
                  # 前i-j-1个卷积层不用变通道，只进行下采样
                  for k in range(i - j - 1):
                      ops.append(
                          nn.Sequential(
                              nn.Conv2d(
                                  c * (2**j),
                                  c * (2**j),
                                  kernel_size=3,
                                  stride=2,
                                  padding=1,
                                  bias=False,
                              ),
                              nn.BatchNorm2d(
                                  c * (2**j), momentum=BN_MOMENTUM
                              ),
                              nn.ReLU(inplace=True),
                          )
                      )
                  # 最后一个卷积层不仅要调整通道，还要进行下采样
                  ops.append(
                      nn.Sequential(
                          nn.Conv2d(
                              c * (2**j),
                              c * (2**i),
                              kernel_size=3,
                              stride=2,
                              padding=1,
                              bias=False,
                          ),
                          nn.BatchNorm2d(
                              c * (2**i), momentum=BN_MOMENTUM
                          ),
                      )
                  )
                  self.fuse_layers[-1].append(nn.Sequential(*ops))
  
      self.relu = nn.ReLU(inplace=True)
  
file_path: deep-learning-for-image-processing/pytorch_keypoint/HRNet/model/hrnet.py
imports: import torch.nn
label_desc: Use image node_id 1 to create a new StageModule object from inherited base classes: nn with example: obj = StageModule(input_branches, output_branches, c)