code: def space_nonlocal(
      input, dim_in, dim_out, prefix, dim_inner, max_pool_stride=2
  ):
      cur = input
      theta = fluid.layers.conv2d(
          input=cur,
          num_filters=dim_inner,
          filter_size=[1, 1],
          stride=[1, 1],
          padding=[0, 0],
          param_attr=ParamAttr(
              name=prefix + "_theta" + "_w",
              initializer=fluid.initializer.Normal(
                  loc=0.0, scale=nonlocal_params["conv_init_std"]
              ),
          ),
          bias_attr=ParamAttr(
              name=prefix + "_theta" + "_b",
              initializer=fluid.initializer.Constant(value=0.0),
          )
          if not nonlocal_params["no_bias"]
          else False,
          name=prefix + "_theta",
      )
      theta_shape = theta.shape
      theta_shape_op = fluid.layers.shape(theta)
      theta_shape_op.stop_gradient = True
  
      if nonlocal_params["use_maxpool"]:
          max_pool = fluid.layers.pool2d(
              input=cur,
              pool_size=[max_pool_stride, max_pool_stride],
              pool_type="max",
              pool_stride=[max_pool_stride, max_pool_stride],
              pool_padding=[0, 0],
              name=prefix + "_pool",
          )
      else:
          max_pool = cur
  
      phi = fluid.layers.conv2d(
          input=max_pool,
          num_filters=dim_inner,
          filter_size=[1, 1],
          stride=[1, 1],
          padding=[0, 0],
          param_attr=ParamAttr(
              name=prefix + "_phi" + "_w",
              initializer=fluid.initializer.Normal(
                  loc=0.0, scale=nonlocal_params["conv_init_std"]
              ),
          ),
          bias_attr=ParamAttr(
              name=prefix + "_phi" + "_b",
              initializer=fluid.initializer.Constant(value=0.0),
          )
          if (nonlocal_params["no_bias"] == 0)
          else False,
          name=prefix + "_phi",
      )
      phi_shape = phi.shape
  
      g = fluid.layers.conv2d(
          input=max_pool,
          num_filters=dim_inner,
          filter_size=[1, 1],
          stride=[1, 1],
          padding=[0, 0],
          param_attr=ParamAttr(
              name=prefix + "_g" + "_w",
              initializer=fluid.initializer.Normal(
                  loc=0.0, scale=nonlocal_params["conv_init_std"]
              ),
          ),
          bias_attr=ParamAttr(
              name=prefix + "_g" + "_b",
              initializer=fluid.initializer.Constant(value=0.0),
          )
          if (nonlocal_params["no_bias"] == 0)
          else False,
          name=prefix + "_g",
      )
      g_shape = g.shape
      # we have to use explicit batch size (to support arbitrary spacetime size)
      # e.g. (8, 1024, 4, 14, 14) => (8, 1024, 784)
      theta = fluid.layers.reshape(theta, shape=(0, 0, -1))
      theta = fluid.layers.transpose(theta, [0, 2, 1])
      phi = fluid.layers.reshape(phi, [0, 0, -1])
      theta_phi = fluid.layers.matmul(
          theta, phi, name=prefix + "_affinity"
      )
      g = fluid.layers.reshape(g, [0, 0, -1])
  
      if nonlocal_params["use_softmax"]:
          if nonlocal_params["use_scale"]:
              theta_phi_sc = fluid.layers.scale(
                  theta_phi, scale=dim_inner**-0.5
              )
          else:
              theta_phi_sc = theta_phi
          p = fluid.layers.softmax(
              theta_phi_sc, name=prefix + "_affinity" + "_prob"
          )
      else:
          # not clear about what is doing in xlw's code
          p = None  # not implemented
          raise "Not implemented when not use softmax"
  
      # note g's axis[2] corresponds to p's axis[2]
      # e.g. g(8, 1024, 784_2) * p(8, 784_1, 784_2) => (8, 1024, 784_1)
      p = fluid.layers.transpose(p, [0, 2, 1])
      t = fluid.layers.matmul(g, p, name=prefix + "_y")
  
      # reshape back
      # e.g. (8, 1024, 784) => (8, 1024, 4, 14, 14)
      t_shape = t.shape
      t_re = fluid.layers.reshape(
          t, shape=list(theta_shape), actual_shape=theta_shape_op
      )
      blob_out = t_re
      blob_out = fluid.layers.conv2d(
          input=blob_out,
          num_filters=dim_out,
          filter_size=[1, 1],
          stride=[1, 1],
          padding=[0, 0],
          param_attr=ParamAttr(
              name=prefix + "_out" + "_w",
              initializer=fluid.initializer.Constant(value=0.0)
              if nonlocal_params["use_zero_init_conv"]
              else fluid.initializer.Normal(
                  loc=0.0, scale=nonlocal_params["conv_init_std"]
              ),
          ),
          bias_attr=ParamAttr(
              name=prefix + "_out" + "_b",
              initializer=fluid.initializer.Constant(value=0.0),
          )
          if (nonlocal_params["no_bias"] == 0)
          else False,
          name=prefix + "_out",
      )
      blob_out_shape = blob_out.shape
  
      if nonlocal_params["use_bn"]:
          bn_name = prefix + "_bn"
          blob_out = fluid.layers.batch_norm(
              blob_out,  # is_test = test_mode, \
              momentum=nonlocal_params["bn_momentum"],
              epsilon=nonlocal_params["bn_epsilon"],
              name=bn_name,
              param_attr=ParamAttr(
                  name=bn_name + "_s",
                  initializer=fluid.initializer.Constant(
                      value=nonlocal_params["bn_init_gamma"]
                  ),
                  regularizer=fluid.regularizer.L2Decay(
                      nonlocal_params["weight_decay_bn"]
                  ),
              ),
              bias_attr=ParamAttr(
                  name=bn_name + "_b",
                  regularizer=fluid.regularizer.L2Decay(
                      nonlocal_params["weight_decay_bn"]
                  ),
              ),
              moving_mean_name=bn_name + "_rm",
              moving_variance_name=bn_name + "_riv",
          )  # add bn
  
      if nonlocal_params["use_affine"]:
          affine_scale = fluid.layers.create_parameter(
              shape=[blob_out_shape[1]],
              dtype=blob_out.dtype,
              attr=ParamAttr(name=prefix + "_affine" + "_s"),
              default_initializer=fluid.initializer.Constant(value=1.0),
          )
          affine_bias = fluid.layers.create_parameter(
              shape=[blob_out_shape[1]],
              dtype=blob_out.dtype,
              attr=ParamAttr(name=prefix + "_affine" + "_b"),
              default_initializer=fluid.initializer.Constant(value=0.0),
          )
          blob_out = fluid.layers.affine_channel(
              blob_out,
              scale=affine_scale,
              bias=affine_bias,
              name=prefix + "_affine",
          )  # add affine
  
      return blob_out
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from __future__ import unicode_literals
  import paddle.fluid
  from paddle.fluid import ParamAttr
label_desc: Use image node_id 1 for calling a global function with example usage: space_nonlocal(input, dim_in, dim_out, prefix, dim_inner, max_pool_stride) and returns: blob_out