parameters: (self, x)
return_statement: return out
docstring: :param x: (n, c, h, w)
docstring_summary: :param x: (n, c, h, w)
docstring_tokens: : param x : ( n c h w )
function: def forward(self, x):
          '''
          :param x: (n, c, h, w)
          '''
          batch_size = x.size(0)
  
          # (n, num_in, h, w) --> (n, num_state, h, w)
          #                   --> (n, num_state, h*w)
          x_state_reshaped = self.conv_state(x).view(batch_size, self.num_s, -1)
  
          # (n, num_in, h, w) --> (n, num_node, h, w)
          #                   --> (n, num_node, h*w)
          x_proj_reshaped = self.conv_proj(x).view(batch_size, self.num_n, -1)
          # x_proj_reshaped = torch.nn.functional.softmax(x_proj_reshaped, dim=2)
  
          # (n, num_in, h, w) --> (n, num_node, h, w)
          #                   --> (n, num_node, h*w)
          x_rproj_reshaped = x_proj_reshaped
  
          # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
  
          # projection: pixel space -> instance space
          # (n, num_state, h*w) x (n, num_node, h*w)T --> (n, num_state, num_node)
          x_n_state = torch.matmul(x_state_reshaped, x_proj_reshaped.permute(0, 2, 1))
          if self.norm_type // 2 == 0:
              x_n_state = x_n_state * (1. / x_state_reshaped.size(2)) # TODO: Grid search
          elif self.norm_type // 2 == 1:
              x_n_state = x_n_state * (1000. / x_state_reshaped.size(2))
          else:
              assert 0
          # x_n_state = self.gcn_bn(x_n_state.view(batch_size, 1, -1)).view(batch_size, self.num_s, self.num_n)
  
          # (n, num_state, num_node) -> (n, num_node, num_state)
          #                          -> (n, num_state, num_node)
          # x_n_rel = self.GCN_G(x_n_state.permute(0, 2, 1).contiguous()).permute(0, 2, 1)
          # x_n_rel = x_n_rel.contiguous() + x_n_state
          # # (n, num_state, num_node) -> (n, num_state, num_node)
          x_n_rel = self.gcn(x_n_state)
  
          # reverse projection: instance space -> pixel space
          # (n, num_state, num_node) x (n, num_node, h*w) --> (n, num_state, h*w)
          x_state_reshaped = torch.matmul(x_n_rel, x_rproj_reshaped)
  
          # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
  
          # (n, num_state, h*w) --> (n, num_state, h, w)
          x_state = x_state_reshaped.view(batch_size, self.num_s, *x.size()[2:])
  
          # -----------------
          # final
          if hasattr(self, 'gamma'):
             out = x + self.gamma * self.fc_2(x_state)
          elif hasattr(self, 'blocker'):
             out = x + self.blocker(self.fc_2(x_state).view(batch_size, 1, -1)).view(x.shape)
          else:
              assert 0
  
          return out
function_tokens: def forward ( self , x ) : batch_size = x . size ( 0 ) # (n, num_in, h, w) --> (n, num_state, h, w) #                   --> (n, num_state, h*w) x_state_reshaped = self . conv_state ( x ) . view ( batch_size , self . num_s , - 1 ) # (n, num_in, h, w) --> (n, num_node, h, w) #                   --> (n, num_node, h*w) x_proj_reshaped = self . conv_proj ( x ) . view ( batch_size , self . num_n , - 1 ) # x_proj_reshaped = torch.nn.functional.softmax(x_proj_reshaped, dim=2) # (n, num_in, h, w) --> (n, num_node, h, w) #                   --> (n, num_node, h*w) x_rproj_reshaped = x_proj_reshaped # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # projection: pixel space -> instance space # (n, num_state, h*w) x (n, num_node, h*w)T --> (n, num_state, num_node) x_n_state = torch . matmul ( x_state_reshaped , x_proj_reshaped . permute ( 0 , 2 , 1 ) ) if self . norm_type // 2 == 0 : x_n_state = x_n_state * ( 1. / x_state_reshaped . size ( 2 ) ) # TODO: Grid search elif self . norm_type // 2 == 1 : x_n_state = x_n_state * ( 1000. / x_state_reshaped . size ( 2 ) ) else : assert 0 # x_n_state = self.gcn_bn(x_n_state.view(batch_size, 1, -1)).view(batch_size, self.num_s, self.num_n) # (n, num_state, num_node) -> (n, num_node, num_state) #                          -> (n, num_state, num_node) # x_n_rel = self.GCN_G(x_n_state.permute(0, 2, 1).contiguous()).permute(0, 2, 1) # x_n_rel = x_n_rel.contiguous() + x_n_state # # (n, num_state, num_node) -> (n, num_state, num_node) x_n_rel = self . gcn ( x_n_state ) # reverse projection: instance space -> pixel space # (n, num_state, num_node) x (n, num_node, h*w) --> (n, num_state, h*w) x_state_reshaped = torch . matmul ( x_n_rel , x_rproj_reshaped ) # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # (n, num_state, h*w) --> (n, num_state, h, w) x_state = x_state_reshaped . view ( batch_size , self . num_s , * x . size ( ) [ 2 : ] ) # ----------------- # final if hasattr ( self , 'gamma' ) : out = x + self . gamma * self . fc_2 ( x_state ) elif hasattr ( self , 'blocker' ) : out = x + self . blocker ( self . fc_2 ( x_state ) . view ( batch_size , 1 , - 1 ) ) . view ( x . shape ) else : assert 0 return out