class_docstr: The basic implementation for pyramid self-attention block/non-local block
class_docstr_tok: ["The","basic","implementation","for","pyramid","self-attention","block\/non-local","block"]
code: def forward(self, x):
      batch_size, c, w, h = x.size()
  
      local_x = []
      local_y = []
      step_w, step_h = w // self.scale, h // self.scale
      for i in range(self.scale):
          for j in range(self.scale):
              start_x, start_y = step_w * i, step_h * j
              end_x, end_y = min(start_x + step_w, w), min(
                  start_y + step_h, h
              )
              if i == (self.scale - 1):
                  end_x = w
              if j == (self.scale - 1):
                  end_y = h
              local_x += [start_x, end_x]
              local_y += [start_y, end_y]
  
      value = self.f_value(x)
      query = self.f_query(x)
      key = self.f_key(x)
  
      local_list = []
      local_block_cnt = (self.scale**2) * 2
      for i in range(0, local_block_cnt, 2):
          value_local = value[
              :,
              :,
              local_x[i] : local_x[i + 1],
              local_y[i] : local_y[i + 1],
          ]
          query_local = query[
              :,
              :,
              local_x[i] : local_x[i + 1],
              local_y[i] : local_y[i + 1],
          ]
          key_local = key[
              :,
              :,
              local_x[i] : local_x[i + 1],
              local_y[i] : local_y[i + 1],
          ]
  
          w_local, h_local = value_local.size(2), value_local.size(3)
          value_local = (
              value_local.contiguous()
              .view(batch_size, self.value_channels, -1)
              .permute(0, 2, 1)
          )
          query_local = (
              query_local.contiguous()
              .view(batch_size, self.key_channels, -1)
              .permute(0, 2, 1)
          )
          key_local = key_local.contiguous().view(
              batch_size, self.key_channels, -1
          )
  
          sim_map = torch.bmm(query_local, key_local) * (
              self.key_channels**-0.5
          )
          sim_map = F.softmax(sim_map, dim=-1)
  
          context_local = (
              torch.bmm(sim_map, value_local)
              .permute(0, 2, 1)
              .contiguous()
          )
          context_local = context_local.view(
              batch_size, self.value_channels, w_local, h_local
          )
          local_list.append(context_local)
  
      context_list = []
      for i in range(0, self.scale):
          row_tmp = []
          for j in range(self.scale):
              row_tmp.append(local_list[j + i * self.scale])
          context_list.append(torch.cat(row_tmp, 3))
  
      context = torch.cat(context_list, 2)
      context = self.W(context)
  
      return context
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  from pywick.models.segmentation.da_basenets.segbase import SegBaseModel
  from pywick.models.segmentation.da_basenets.fcn import _FCNHead
label_desc: Use image node_id 2 for calling the PyramidAttentionBlock obj's underlying member method code with example usage: obj.forward(x) and returns: context