code: import torch.nn as nn
  import torch.nn.functional as F
  import torch
  
  
  class CPM(nn.Module):
      def __init__(self, k):
          super(CPM, self).__init__()
          self.k = k
          self.pool_center = nn.AvgPool2d(kernel_size=9, stride=8, padding=1)
          self.conv1_stage1 = nn.Conv2d(3, 128, kernel_size=9, padding=4)
          self.pool1_stage1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
          self.conv2_stage1 = nn.Conv2d(128, 128, kernel_size=9, padding=4)
          self.pool2_stage1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
          self.conv3_stage1 = nn.Conv2d(128, 128, kernel_size=9, padding=4)
          self.pool3_stage1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
          self.conv4_stage1 = nn.Conv2d(128, 32, kernel_size=5, padding=2)
          self.conv5_stage1 = nn.Conv2d(32, 512, kernel_size=9, padding=4)
          self.conv6_stage1 = nn.Conv2d(512, 512, kernel_size=1)
          self.conv7_stage1 = nn.Conv2d(512, self.k + 1, kernel_size=1)
  
          self.conv1_stage2 = nn.Conv2d(3, 128, kernel_size=9, padding=4)
          self.pool1_stage2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
          self.conv2_stage2 = nn.Conv2d(128, 128, kernel_size=9, padding=4)
          self.pool2_stage2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
          self.conv3_stage2 = nn.Conv2d(128, 128, kernel_size=9, padding=4)
          self.pool3_stage2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
          self.conv4_stage2 = nn.Conv2d(128, 32, kernel_size=5, padding=2)
  
          self.Mconv1_stage2 = nn.Conv2d(32 + self.k + 2, 128, kernel_size=11, padding=5)
          self.Mconv2_stage2 = nn.Conv2d(128, 128, kernel_size=11, padding=5)
          self.Mconv3_stage2 = nn.Conv2d(128, 128, kernel_size=11, padding=5)
          self.Mconv4_stage2 = nn.Conv2d(128, 128, kernel_size=1, padding=0)
          self.Mconv5_stage2 = nn.Conv2d(128, self.k + 1, kernel_size=1, padding=0)
  
          self.conv1_stage3 = nn.Conv2d(128, 32, kernel_size=5, padding=2)
  
          self.Mconv1_stage3 = nn.Conv2d(32 + self.k + 2, 128, kernel_size=11, padding=5)
          self.Mconv2_stage3 = nn.Conv2d(128, 128, kernel_size=11, padding=5)
          self.Mconv3_stage3 = nn.Conv2d(128, 128, kernel_size=11, padding=5)
          self.Mconv4_stage3 = nn.Conv2d(128, 128, kernel_size=1, padding=0)
          self.Mconv5_stage3 = nn.Conv2d(128, self.k + 1, kernel_size=1, padding=0)
  
          self.conv1_stage4 = nn.Conv2d(128, 32, kernel_size=5, padding=2)
  
          self.Mconv1_stage4 = nn.Conv2d(32 + self.k + 2, 128, kernel_size=11, padding=5)
          self.Mconv2_stage4 = nn.Conv2d(128, 128, kernel_size=11, padding=5)
          self.Mconv3_stage4 = nn.Conv2d(128, 128, kernel_size=11, padding=5)
          self.Mconv4_stage4 = nn.Conv2d(128, 128, kernel_size=1, padding=0)
          self.Mconv5_stage4 = nn.Conv2d(128, self.k + 1, kernel_size=1, padding=0)
  
          self.conv1_stage5 = nn.Conv2d(128, 32, kernel_size=5, padding=2)
  
          self.Mconv1_stage5 = nn.Conv2d(32 + self.k + 2, 128, kernel_size=11, padding=5)
          self.Mconv2_stage5 = nn.Conv2d(128, 128, kernel_size=11, padding=5)
          self.Mconv3_stage5 = nn.Conv2d(128, 128, kernel_size=11, padding=5)
          self.Mconv4_stage5 = nn.Conv2d(128, 128, kernel_size=1, padding=0)
          self.Mconv5_stage5 = nn.Conv2d(128, self.k + 1, kernel_size=1, padding=0)
  
          self.conv1_stage6 = nn.Conv2d(128, 32, kernel_size=5, padding=2)
  
          self.Mconv1_stage6 = nn.Conv2d(32 + self.k + 2, 128, kernel_size=11, padding=5)
          self.Mconv2_stage6 = nn.Conv2d(128, 128, kernel_size=11, padding=5)
          self.Mconv3_stage6 = nn.Conv2d(128, 128, kernel_size=11, padding=5)
          self.Mconv4_stage6 = nn.Conv2d(128, 128, kernel_size=1, padding=0)
          self.Mconv5_stage6 = nn.Conv2d(128, self.k + 1, kernel_size=1, padding=0)
  
      def _stage1(self, image):
          x = self.pool1_stage1(F.relu(self.conv1_stage1(image)))
          x = self.pool2_stage1(F.relu(self.conv2_stage1(x)))
          x = self.pool3_stage1(F.relu(self.conv3_stage1(x)))
          x = F.relu(self.conv4_stage1(x))
          x = F.relu(self.conv5_stage1(x))
          x = F.relu(self.conv6_stage1(x))
          x = self.conv7_stage1(x)
  
          return x
  
      def _middle(self, image):
          x = self.pool1_stage2(F.relu(self.conv1_stage2(image)))
          x = self.pool2_stage2(F.relu(self.conv2_stage2(x)))
          x = self.pool3_stage2(F.relu(self.conv3_stage2(x)))
  
          return x
  
      def _stage2(self, pool3_stage2_map, conv7_stage1_map, pool_center_map):
          x = F.relu(self.conv4_stage2(pool3_stage2_map))
          x = torch.cat([x, conv7_stage1_map, pool_center_map], dim=1)
          x = F.relu(self.Mconv1_stage2(x))
          x = F.relu(self.Mconv2_stage2(x))
          x = F.relu(self.Mconv3_stage2(x))
          x = F.relu(self.Mconv4_stage2(x))
          x = self.Mconv5_stage2(x)
  
          return x
  
      def _stage3(self, pool3_stage2_map, Mconv5_stage2_map, pool_center_map):
          x = F.relu(self.conv1_stage3(pool3_stage2_map))
          x = torch.cat([x, Mconv5_stage2_map, pool_center_map], dim=1)
          x = F.relu(self.Mconv1_stage3(x))
          x = F.relu(self.Mconv2_stage3(x))
          x = F.relu(self.Mconv3_stage3(x))
          x = F.relu(self.Mconv4_stage3(x))
          x = self.Mconv5_stage3(x)
  
          return x
  
      def _stage4(self, pool3_stage2_map, Mconv5_stage3_map, pool_center_map):
          x = F.relu(self.conv1_stage4(pool3_stage2_map))
          x = torch.cat([x, Mconv5_stage3_map, pool_center_map], dim=1)
          x = F.relu(self.Mconv1_stage4(x))
          x = F.relu(self.Mconv2_stage4(x))
          x = F.relu(self.Mconv3_stage4(x))
          x = F.relu(self.Mconv4_stage4(x))
          x = self.Mconv5_stage4(x)
  
          return x
  
      def _stage5(self, pool3_stage2_map, Mconv5_stage4_map, pool_center_map):
          x = F.relu(self.conv1_stage5(pool3_stage2_map))
          x = torch.cat([x, Mconv5_stage4_map, pool_center_map], dim=1)
          x = F.relu(self.Mconv1_stage5(x))
          x = F.relu(self.Mconv2_stage5(x))
          x = F.relu(self.Mconv3_stage5(x))
          x = F.relu(self.Mconv4_stage5(x))
          x = self.Mconv5_stage5(x)
  
          return x
  
      def _stage6(self, pool3_stage2_map, Mconv5_stage5_map, pool_center_map):
          x = F.relu(self.conv1_stage6(pool3_stage2_map))
          x = torch.cat([x, Mconv5_stage5_map, pool_center_map], dim=1)
          x = F.relu(self.Mconv1_stage6(x))
          x = F.relu(self.Mconv2_stage6(x))
          x = F.relu(self.Mconv3_stage6(x))
          x = F.relu(self.Mconv4_stage6(x))
          x = self.Mconv5_stage6(x)
  
          return x
  
      def forward(self, image, center_map):
          pool_center_map = self.pool_center(center_map)
  
          conv7_stage1_map = self._stage1(image)
  
          pool3_stage2_map = self._middle(image)
  
          Mconv5_stage2_map = self._stage2(pool3_stage2_map, conv7_stage1_map,
                                           pool_center_map)
          Mconv5_stage3_map = self._stage3(pool3_stage2_map, Mconv5_stage2_map,
                                           pool_center_map)
          Mconv5_stage4_map = self._stage4(pool3_stage2_map, Mconv5_stage3_map,
                                           pool_center_map)
          Mconv5_stage5_map = self._stage5(pool3_stage2_map, Mconv5_stage4_map,
                                           pool_center_map)
          Mconv5_stage6_map = self._stage6(pool3_stage2_map, Mconv5_stage5_map,
                                           pool_center_map)
  
          return conv7_stage1_map, Mconv5_stage2_map, Mconv5_stage3_map, Mconv5_stage4_map, Mconv5_stage5_map, Mconv5_stage6_map
  
  
  def get_pose_net(cfg, is_train, **kwargs):
      model = CPM(16)
  
      return model
  
  
  if __name__ == '__main__':
      from torchsummary import summary
  
      model = get_pose_net('', True).cuda()
      summary(model, (3, 256, 256))
apis: torch.nn.MaxPool2d torch.nn.Conv2d torch.cat torch.nn.AvgPool2d