parameters: (main_program=None,
                  n=2,
                  m=4,
                  mask_algo='mask_1d',
                  with_mask=True,
                  sharding=False)
return_statement: return ASPHelper.prune_model(
          place=place,
          main_program=main_program,
          n=n,
          m=m,
          mask_algo=MaskAlgo_mapping[mask_algo],
          with_mask=with_mask)
docstring: r"""
      Pruning parameters of supported layers in :attr:`main_program` via 
      specified mask generation function given by :attr:`mask_algo`. This 
      function supports both training and inference controlled by :attr:`with_mask`.
      If :attr:`with_mask` is True, it would also prune parameter related ASP mask Variables,
      else only prunes parameters.
  
      *Note*: If parameters are supported and in FP16, please set :attr:`n`=2, :attr:`m`=4, 
      if they in FP32, then :attr:`n`=1, :attr:`m`=2` to further enable Sparse Tensor Core acceleration.
  
      *Note*: If calling this function with :attr:`with_mask`, it should call `OptimizerWithSparsityGuarantee.minimize` 
      and initialization (`exe.run(startup_program`)) before (For successfully obtain mask Variable). 
      Typically set `with_mask` as true for training (have called `OptimizerWithSparsityGuarantee.minimize`) and false for 
      inference only. To obtain OptimizerWithSparsityGuarantee, please see `sparsity.decoreate()`.
  
      Args:
          main_program (Program, optional): Program with model definition and its parameters. Default is `paddle.static.default_main_program()
          n (int): n of `n:m` sparse pattern.
          m (int): m of `n:m` sparse pattern.
          mask_algo (string, optional): The function name to generate spase mask. Default is `mask_1d`.
                                        The vaild inputs should be one of 'mask_1d', 'mask_2d_greedy' and 'mask_2d_best'.
          with_mask (bool, optional): To prune mask Variables related to parameters or not. Ture is purning also, False is not. Defalut is True.
          sharding (bool, optional): Whether to turn on sharding (model parallel) during training. Please consider turning it ON when encountering OOM using sharding. Default is False.
      Returns:
          dictionary: A dictionary with key: `parameter name` (string) and value: its corresponding mask Variable.
      Examples:
          .. code-block:: python
  
              import paddle
              from paddle.static import sparsity
  
              paddle.enable_static()
  
              main_program = paddle.static.Program()
              startup_program = paddle.static.Program()
  
              with paddle.static.program_guard(main_program, startup_program):
                  input_data = paddle.static.data(name='data', shape=[None, 128])
                  label = paddle.static.data(name='label', shape=[None, 10])
                  hidden = paddle.static.nn.fc(x=input_data, num_flatten_dims=-1, size=32, activation=None, name="need_sparse_fc")
                  hidden = paddle.static.nn.fc(x=hidden, num_flatten_dims=-1, size=32, activation=None, name="need_dense_fc")
                  prob = paddle.static.nn.fc(x=hidden, num_flatten_dims=-1, size=10, activation=None)
                  loss = paddle.mean(paddle.nn.functional.square_error_cost(prob, label))
  
                  # Setup exluded layers out from ASP workflow.
                  # Please note, excluded_layers must be set before calling `optimizer.minimize()`.
                  sparsity.set_excluded_layers(main_program, ["need_dense_fc"])
  
                  optimizer = paddle.optimizer.SGD(learning_rate=0.1)
                  optimizer = paddle.static.amp.decorate(optimizer )
                  # Calling sparsity.decorate() to wrap minimize() in optimizer, which 
                  # will insert necessary masking operations for ASP workflow.
                  optimizer = sparsity.decorate(optimizer)
                  optimizer.minimize(loss, startup_program)
  
              device = paddle.device.get_device()
              place = paddle.set_device(device)
  
              exe = paddle.static.Executor(place)
              exe.run(startup_program)
  
              # Must call `exe.run(startup_program)` first before calling `sparsity.prune_model`
              sparsity.prune_model(main_program, mask_algo='mask_2d_best')
docstring_summary: r"""
      Pruning parameters of supported layers in :attr:`main_program` via 
      specified mask generation function given by :attr:`mask_algo`. This 
      function supports both training and inference controlled by :attr:`with_mask`.
      If :attr:`with_mask` is True, it would also prune parameter related ASP mask Variables,
      else only prunes parameters.
docstring_tokens: r Pruning parameters of supported layers in : attr : main_program via specified mask generation function given by : attr : mask_algo . This function supports both training and inference controlled by : attr : with_mask . If : attr : with_mask is True it would also prune parameter related ASP mask Variables else only prunes parameters .
function: def prune_model(main_program=None,
                  n=2,
                  m=4,
                  mask_algo='mask_1d',
                  with_mask=True,
                  sharding=False):
      r"""
      Pruning parameters of supported layers in :attr:`main_program` via 
      specified mask generation function given by :attr:`mask_algo`. This 
      function supports both training and inference controlled by :attr:`with_mask`.
      If :attr:`with_mask` is True, it would also prune parameter related ASP mask Variables,
      else only prunes parameters.
  
      *Note*: If parameters are supported and in FP16, please set :attr:`n`=2, :attr:`m`=4, 
      if they in FP32, then :attr:`n`=1, :attr:`m`=2` to further enable Sparse Tensor Core acceleration.
  
      *Note*: If calling this function with :attr:`with_mask`, it should call `OptimizerWithSparsityGuarantee.minimize` 
      and initialization (`exe.run(startup_program`)) before (For successfully obtain mask Variable). 
      Typically set `with_mask` as true for training (have called `OptimizerWithSparsityGuarantee.minimize`) and false for 
      inference only. To obtain OptimizerWithSparsityGuarantee, please see `sparsity.decoreate()`.
  
      Args:
          main_program (Program, optional): Program with model definition and its parameters. Default is `paddle.static.default_main_program()
          n (int): n of `n:m` sparse pattern.
          m (int): m of `n:m` sparse pattern.
          mask_algo (string, optional): The function name to generate spase mask. Default is `mask_1d`.
                                        The vaild inputs should be one of 'mask_1d', 'mask_2d_greedy' and 'mask_2d_best'.
          with_mask (bool, optional): To prune mask Variables related to parameters or not. Ture is purning also, False is not. Defalut is True.
          sharding (bool, optional): Whether to turn on sharding (model parallel) during training. Please consider turning it ON when encountering OOM using sharding. Default is False.
      Returns:
          dictionary: A dictionary with key: `parameter name` (string) and value: its corresponding mask Variable.
      Examples:
          .. code-block:: python
  
              import paddle
              from paddle.static import sparsity
  
              paddle.enable_static()
  
              main_program = paddle.static.Program()
              startup_program = paddle.static.Program()
  
              with paddle.static.program_guard(main_program, startup_program):
                  input_data = paddle.static.data(name='data', shape=[None, 128])
                  label = paddle.static.data(name='label', shape=[None, 10])
                  hidden = paddle.static.nn.fc(x=input_data, num_flatten_dims=-1, size=32, activation=None, name="need_sparse_fc")
                  hidden = paddle.static.nn.fc(x=hidden, num_flatten_dims=-1, size=32, activation=None, name="need_dense_fc")
                  prob = paddle.static.nn.fc(x=hidden, num_flatten_dims=-1, size=10, activation=None)
                  loss = paddle.mean(paddle.nn.functional.square_error_cost(prob, label))
  
                  # Setup exluded layers out from ASP workflow.
                  # Please note, excluded_layers must be set before calling `optimizer.minimize()`.
                  sparsity.set_excluded_layers(main_program, ["need_dense_fc"])
  
                  optimizer = paddle.optimizer.SGD(learning_rate=0.1)
                  optimizer = paddle.static.amp.decorate(optimizer )
                  # Calling sparsity.decorate() to wrap minimize() in optimizer, which 
                  # will insert necessary masking operations for ASP workflow.
                  optimizer = sparsity.decorate(optimizer)
                  optimizer.minimize(loss, startup_program)
  
              device = paddle.device.get_device()
              place = paddle.set_device(device)
  
              exe = paddle.static.Executor(place)
              exe.run(startup_program)
  
              # Must call `exe.run(startup_program)` first before calling `sparsity.prune_model`
              sparsity.prune_model(main_program, mask_algo='mask_2d_best')
      """
      if sharding:
          gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0))
          place = paddle.CUDAPlace(gpu_id)
      else:
          device = paddle.device.get_device()
          place = paddle.set_device(device)
  
      MaskAlgo_mapping = {
          'mask_1d': sparsity.MaskAlgo.MASK_1D,
          'mask_2d_greedy': sparsity.MaskAlgo.MASK_2D_GREEDY,
          'mask_2d_best': sparsity.MaskAlgo.MASK_2D_BEST
      }
      assert (mask_algo in MaskAlgo_mapping), \
             'The "mask_algo" should be one of ["mask_1d", "mask_2d_greedy", "mask_2d_best"]'
  
      return ASPHelper.prune_model(
          place=place,
          main_program=main_program,
          n=n,
          m=m,
          mask_algo=MaskAlgo_mapping[mask_algo],
          with_mask=with_mask)
function_tokens: def prune_model ( main_program = None , n = 2 , m = 4 , mask_algo = 'mask_1d' , with_mask = True , sharding = False ) : if sharding : gpu_id = int ( os . environ . get ( 'FLAGS_selected_gpus' , 0 ) ) place = paddle . CUDAPlace ( gpu_id ) else : device = paddle . device . get_device ( ) place = paddle . set_device ( device ) MaskAlgo_mapping = { 'mask_1d' : sparsity . MaskAlgo . MASK_1D , 'mask_2d_greedy' : sparsity . MaskAlgo . MASK_2D_GREEDY , 'mask_2d_best' : sparsity . MaskAlgo . MASK_2D_BEST } assert ( mask_algo in MaskAlgo_mapping ) , 'The "mask_algo" should be one of ["mask_1d", "mask_2d_greedy", "mask_2d_best"]' return ASPHelper . prune_model ( place = place , main_program = main_program , n = n , m = m , mask_algo = MaskAlgo_mapping [ mask_algo ] , with_mask = with_mask )