decorate¶

paddle.static.sparsity. decorate ( optimizer ) [source]

Wrap the given optimizer as a OptimizerWithSparsityGuarantee, which would insert necessary ops for ASP workflows when calling minimize()

Parameters

optimizer (Optimizer) – A Optimizer used for training.

Returns

A wrapper for ASP to decorate minimize function of the given optimizer.

Return type

OptimizerWithSparsityGuarantee

Examples

import paddle
from paddle.static import sparsity

main_program = paddle.static.Program()
startup_program = paddle.static.Program()

paddle.enable_static()

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)
    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))

    optimizer = paddle.optimizer.SGD(learning_rate=0.1)
    optimizer = sparsity.decorate(optimizer)
    # if do sparse training with Fleet, please replace above decorate with:
    # strategy = paddle.distributed.fleet.DistributedStrategy()
    # strategy.asp = True
    # optimizer = fleet.distributed_optimizer(optimizer, strategy=strategy)

    optimizer.minimize(loss, startup_program)