wehub-resource-sync
1100 lines
46 KiB
Python
1100 lines
46 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
# Copyright (c) 2021 NVIDIA Corporation. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
Functions for Auto SParsity (ASP) training and inference.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import os
|
|
from typing import TYPE_CHECKING
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle.base import core, global_scope, program_guard
|
|
from paddle.base.framework import dygraph_only
|
|
|
|
from .supported_layer_list import (
|
|
_default_pruning,
|
|
supported_layers_and_prune_func_map,
|
|
)
|
|
from .utils import MaskAlgo
|
|
|
|
OpRole = core.op_proto_and_checker_maker.OpRole
|
|
OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable, Iterable, Sequence
|
|
from typing import Any, Literal
|
|
|
|
import numpy.typing as npt
|
|
|
|
from paddle import Tensor
|
|
from paddle._typing import PlaceLike
|
|
from paddle.nn import Layer
|
|
from paddle.optimizer import Optimizer
|
|
from paddle.static import Operator, Program
|
|
|
|
__all__ = []
|
|
|
|
|
|
def set_excluded_layers(
|
|
param_names: list[str], main_program: Program | None = None
|
|
) -> None:
|
|
r"""
|
|
Set parameter name of layers which would not be pruned as sparse weights.
|
|
|
|
Args:
|
|
param_names (list of string): A list contains names of parameters.
|
|
main_program (Program|None, optional): Program with model definition and its parameters.
|
|
If None is given, then it would be set as `paddle.static.default_main_program().
|
|
Default is None.
|
|
Examples:
|
|
.. code-block:: pycon
|
|
:name: dynamic-graph
|
|
|
|
>>> # doctest: +SKIP("paddle.incubate.asp module doesn't support PIR mode now")
|
|
>>> # Example1: Usage of Dynamic Graph
|
|
>>> import paddle
|
|
|
|
>>> class MyLayer(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self.conv1 = paddle.nn.Conv2D(
|
|
... in_channels=3,
|
|
... out_channels=4,
|
|
... kernel_size=3,
|
|
... padding=2,
|
|
... )
|
|
... self.linear1 = paddle.nn.Linear(4624, 100)
|
|
...
|
|
... def forward(self, img):
|
|
... hidden = self.conv1(img)
|
|
... hidden = paddle.flatten(hidden, start_axis=1)
|
|
... prediction = self.linear1(hidden)
|
|
... return prediction
|
|
|
|
>>> my_layer = MyLayer()
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=my_layer.parameters(),
|
|
... )
|
|
|
|
>>> # Need to set excluded layers before calling decorate
|
|
>>> paddle.incubate.asp.set_excluded_layers([my_layer.linear1.full_name()])
|
|
|
|
>>> optimizer = paddle.incubate.asp.decorate(optimizer)
|
|
|
|
.. code-block:: pycon
|
|
:name: static-graph
|
|
|
|
>>> # doctest: +SKIP("paddle.incubate.asp module doesn't support PIR mode now")
|
|
>>> # Example2: Usage of Static Graph
|
|
>>> import paddle
|
|
|
|
>>> paddle.enable_static()
|
|
|
|
>>> class MyLayer(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self.conv1 = paddle.nn.Conv2D(
|
|
... in_channels=3,
|
|
... out_channels=4,
|
|
... kernel_size=3,
|
|
... padding=2,
|
|
... )
|
|
... self.linear1 = paddle.nn.Linear(4624, 100)
|
|
...
|
|
... def forward(self, img):
|
|
... hidden = self.conv1(img)
|
|
... hidden = paddle.flatten(hidden, start_axis=1)
|
|
... prediction = self.linear1(hidden)
|
|
... return prediction
|
|
|
|
>>> 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, 3, 224, 224])
|
|
... label = paddle.static.data(name='label', shape=[None, 100])
|
|
... my_layer = MyLayer()
|
|
... prob = my_layer(input_data)
|
|
... loss = paddle.mean(paddle.nn.functional.square_error_cost(prob, label))
|
|
...
|
|
... # Setup excluded layers out from ASP workflow.
|
|
... # Please note, excluded_layers must be set before calling optimizer.minimize().
|
|
... paddle.incubate.asp.set_excluded_layers([my_layer.linear1.full_name()], main_program)
|
|
...
|
|
... optimizer = paddle.optimizer.SGD(learning_rate=0.1)
|
|
... optimizer = paddle.static.amp.decorate(optimizer)
|
|
... # Calling paddle.incubate.asp.decorate() to wrap minimize() in optimizer, which
|
|
... # will insert necessary masking operations for ASP workflow.
|
|
... optimizer = paddle.incubate.asp.decorate(optimizer)
|
|
... optimizer.minimize(loss, startup_program)
|
|
"""
|
|
if main_program is None:
|
|
main_program = paddle.static.default_main_program()
|
|
ASPHelper.set_excluded_layers(
|
|
param_names=param_names, main_program=main_program
|
|
)
|
|
|
|
|
|
def reset_excluded_layers(main_program: Program | None = None) -> None:
|
|
r"""
|
|
Reset excluded layers setting corresponding to :attr:`main_program`. If :attr:`main_program`
|
|
is None, then all configurations of excluded_layers would be cleaned.
|
|
|
|
Args:
|
|
main_program (Program, optional): Program with model definition and its parameters.
|
|
If None is given, then this function would reset all excluded_layers.
|
|
Default is None.
|
|
Examples:
|
|
.. code-block:: pycon
|
|
:name: dynamic-graph
|
|
|
|
>>> # doctest: +SKIP("paddle.incubate.asp module doesn't support PIR mode now")
|
|
>>> # Example1: Usage of Dynamic Graph
|
|
>>> import paddle
|
|
|
|
>>> class MyLayer(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self.conv1 = paddle.nn.Conv2D(
|
|
... in_channels=3,
|
|
... out_channels=4,
|
|
... kernel_size=3,
|
|
... padding=2,
|
|
... )
|
|
... self.linear1 = paddle.nn.Linear(4624, 100)
|
|
...
|
|
... def forward(self, img):
|
|
... hidden = self.conv1(img)
|
|
... hidden = paddle.flatten(hidden, start_axis=1)
|
|
... prediction = self.linear1(hidden)
|
|
... return prediction
|
|
|
|
>>> my_layer = MyLayer()
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=my_layer.parameters(),
|
|
... )
|
|
|
|
>>> # Need to set excluded layers before calling decorate
|
|
>>> paddle.incubate.asp.set_excluded_layers([my_layer.linear1.full_name()])
|
|
>>> # Reset excluded_layers, all supported layers would be included into Automatic SParsity's workflow.
|
|
>>> # Please note, reset_excluded_layers also must be called before calling asp.decorate().
|
|
>>> paddle.incubate.asp.reset_excluded_layers()
|
|
|
|
>>> optimizer = paddle.incubate.asp.decorate(optimizer)
|
|
|
|
.. code-block:: pycon
|
|
:name: static-graph
|
|
|
|
>>> # doctest: +SKIP("paddle.incubate.asp module doesn't support PIR mode now")
|
|
>>> # Example2: Usage of Static Graph
|
|
>>> import paddle
|
|
|
|
>>> paddle.enable_static()
|
|
|
|
>>> class MyLayer(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self.conv1 = paddle.nn.Conv2D(
|
|
... in_channels=3,
|
|
... out_channels=4,
|
|
... kernel_size=3,
|
|
... padding=2,
|
|
... )
|
|
... self.linear1 = paddle.nn.Linear(4624, 100)
|
|
...
|
|
... def forward(self, img):
|
|
... hidden = self.conv1(img)
|
|
... hidden = paddle.flatten(hidden, start_axis=1)
|
|
... prediction = self.linear1(hidden)
|
|
... return prediction
|
|
|
|
>>> 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, 3, 224, 224])
|
|
... label = paddle.static.data(name='label', shape=[None, 100])
|
|
... my_layer = MyLayer()
|
|
... prob = my_layer(input_data)
|
|
... loss = paddle.mean(paddle.nn.functional.square_error_cost(prob, label))
|
|
...
|
|
... # Setup excluded layers out from ASP workflow.
|
|
... # Please note, excluded_layers must be set before calling optimizer.minimize().
|
|
... paddle.incubate.asp.set_excluded_layers([my_layer.linear1.full_name()], main_program)
|
|
... # Reset excluded_layers, all supported layers would be included into Automatic SParsity's workflow.
|
|
... # Please note, reset_excluded_layers also must be called before calling optimizer.minimize().
|
|
... paddle.incubate.asp.reset_excluded_layers(main_program)
|
|
...
|
|
... optimizer = paddle.optimizer.SGD(learning_rate=0.1)
|
|
... optimizer = paddle.static.amp.decorate(optimizer)
|
|
... # Calling paddle.incubate.asp.decorate() to wrap minimize() in optimizer, which
|
|
... # will insert necessary masking operations for ASP workflow.
|
|
... optimizer = paddle.incubate.asp.decorate(optimizer)
|
|
... optimizer.minimize(loss, startup_program)
|
|
"""
|
|
ASPHelper.reset_excluded_layers(main_program=main_program)
|
|
|
|
|
|
def decorate(optimizer: Optimizer) -> OptimizerWithSparsityGuarantee:
|
|
r"""
|
|
Wrap the given optimizer as a OptimizerWithSparsityGuarantee,
|
|
If running with dynamic graph mode. ASP would creates mask variables for supported parameters.
|
|
Else if in static graph mode, ASP would creates mask variables and inserts necessary ops
|
|
when calling minimize()
|
|
|
|
Args:
|
|
optimizer (Optimizer): A Optimizer used for training.
|
|
Returns:
|
|
OptimizerWithSparsityGuarantee: A wrapper for ASP to decorate `minimize` function of the given optimizer.
|
|
Examples:
|
|
.. code-block:: pycon
|
|
:name: dynamic-graph
|
|
|
|
>>> # doctest: +SKIP("paddle.incubate.asp module doesn't support PIR mode now")
|
|
>>> # Example1: Usage of Dynamic Graph
|
|
>>> import paddle
|
|
|
|
>>> class MyLayer(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self.conv1 = paddle.nn.Conv2D(
|
|
... in_channels=3,
|
|
... out_channels=4,
|
|
... kernel_size=3,
|
|
... padding=2,
|
|
... )
|
|
... self.linear1 = paddle.nn.Linear(4624, 32)
|
|
... self.linear2 = paddle.nn.Linear(32, 32)
|
|
... self.linear3 = paddle.nn.Linear(32, 10)
|
|
...
|
|
... def forward(self, img):
|
|
... hidden = self.conv1(img)
|
|
... hidden = paddle.flatten(hidden, start_axis=1)
|
|
... hidden = self.linear1(hidden)
|
|
... hidden = self.linear2(hidden)
|
|
... prediction = self.linear3(hidden)
|
|
... return prediction
|
|
|
|
>>> my_layer = MyLayer()
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=my_layer.parameters(),
|
|
... )
|
|
|
|
>>> # Calling paddle.incubate.asp.decorate() to wrap step() in optimizer, which
|
|
>>> # will apply necessary masking operations for ASP workflow.
|
|
>>> # In dynamic graph mode, ASP would create related mask variables during decoration.
|
|
>>> optimizer = paddle.incubate.asp.decorate(optimizer)
|
|
|
|
.. code-block:: pycon
|
|
:name: static-graph
|
|
|
|
>>> # doctest: +SKIP("paddle.incubate.asp module doesn't support PIR mode now")
|
|
>>> # Example2: Usage of Static Graph
|
|
>>> import paddle
|
|
|
|
>>> paddle.enable_static()
|
|
|
|
>>> class MyLayer(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self.conv1 = paddle.nn.Conv2D(
|
|
... in_channels=3,
|
|
... out_channels=4,
|
|
... kernel_size=3,
|
|
... padding=2,
|
|
... )
|
|
... self.linear1 = paddle.nn.Linear(4624, 100)
|
|
...
|
|
... def forward(self, img):
|
|
... hidden = self.conv1(img)
|
|
... hidden = paddle.flatten(hidden, start_axis=1)
|
|
... prediction = self.linear1(hidden)
|
|
... return prediction
|
|
|
|
>>> 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, 3, 224, 224])
|
|
... label = paddle.static.data(name='label', shape=[None, 100])
|
|
... my_layer = MyLayer()
|
|
... prob = my_layer(input_data)
|
|
... loss = paddle.mean(paddle.nn.functional.square_error_cost(prob, label))
|
|
...
|
|
... optimizer = paddle.optimizer.SGD(learning_rate=0.1)
|
|
... # Calling paddle.incubate.asp.decorate() to wrap minimize() in optimizer, which
|
|
... # will insert necessary masking operations for ASP workflow.
|
|
... # In static graph mode, ASP creates related mask variables
|
|
... # during minimize().
|
|
... optimizer = paddle.incubate.asp.decorate(optimizer)
|
|
... optimizer.minimize(loss, startup_program)
|
|
"""
|
|
return ASPHelper.decorate(optimizer)
|
|
|
|
|
|
def prune_model(
|
|
model: Program | Layer,
|
|
n: int = 2,
|
|
m: int = 4,
|
|
mask_algo: Literal['mask_1d', 'mask_2d_greedy', 'mask_2d_best'] = 'mask_1d',
|
|
with_mask: bool = True,
|
|
) -> dict[str, Tensor]:
|
|
r"""
|
|
Pruning parameters of supported layers in :attr:`model` 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*: (Static graph mode) 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 `paddle.incubate.asp.decorate()`.
|
|
|
|
Args:
|
|
model (Program|nn.Layer): Program with model definition and its parameters, or a object of `paddle.nn.Layer`.
|
|
n (int, optional): n of `n:m` sparse pattern. Default is 2.
|
|
m (int, optional): m of `n:m` sparse pattern. Default is 4.
|
|
mask_algo (string, optional): The function name to generate sparse mask. Default is `mask_1d`.
|
|
The valid 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. True is pruning also, False is not. Default is True.
|
|
Returns:
|
|
dictionary: A dictionary with key: `parameter name` (string) and value: its corresponding mask Variable.
|
|
Examples:
|
|
.. code-block:: pycon
|
|
:name: dynamic-graph
|
|
|
|
>>> # doctest: +SKIP("paddle.incubate.asp module doesn't support PIR mode now")
|
|
>>> # Example1: Usage of Dynamic Graph
|
|
>>> import paddle
|
|
>>> import numpy as np
|
|
|
|
>>> class MyLayer(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self.conv1 = paddle.nn.Conv2D(
|
|
... in_channels=3,
|
|
... out_channels=4,
|
|
... kernel_size=3,
|
|
... padding=2,
|
|
... )
|
|
... self.linear1 = paddle.nn.Linear(4624, 32)
|
|
... self.linear2 = paddle.nn.Linear(32, 32)
|
|
... self.linear3 = paddle.nn.Linear(32, 10)
|
|
...
|
|
... def forward(self, img):
|
|
... hidden = self.conv1(img)
|
|
... hidden = paddle.flatten(hidden, start_axis=1)
|
|
... hidden = self.linear1(hidden)
|
|
... hidden = self.linear2(hidden)
|
|
... prediction = self.linear3(hidden)
|
|
... return prediction
|
|
|
|
>>> my_layer = MyLayer()
|
|
>>> loss_fn = paddle.nn.MSELoss(reduction='mean')
|
|
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=my_layer.parameters(),
|
|
... )
|
|
|
|
>>> # Calling paddle.incubate.asp.decorate() to wrap step() in optimizer, which
|
|
>>> # will apply necessary masking operations for ASP workflow.
|
|
>>> # In dynamic graph mode, ASP would create related mask variables during decoration.
|
|
>>> optimizer = paddle.incubate.asp.decorate(optimizer)
|
|
|
|
>>> # Must call paddle.incubate.asp.decorate() first before calling paddle.incubate.asp.prune_model()
|
|
>>> paddle.incubate.asp.prune_model(my_layer, mask_algo='mask_2d_best')
|
|
|
|
>>> for i in range(10):
|
|
... imgs = paddle.to_tensor(
|
|
... np.random.randn(64, 3, 32, 32),
|
|
... dtype='float32',
|
|
... stop_gradient=False,
|
|
... )
|
|
... labels = paddle.to_tensor(
|
|
... np.random.randint(10, size=(64, 1)),
|
|
... dtype='float32',
|
|
... stop_gradient=False,
|
|
... )
|
|
... output = my_layer(imgs)
|
|
... loss = loss_fn(output, labels)
|
|
... loss.backward()
|
|
... optimizer.step()
|
|
... optimizer.clear_grad()
|
|
|
|
.. code-block:: pycon
|
|
:name: static-graph
|
|
|
|
>>> # doctest: +SKIP("paddle.incubate.asp module doesn't support PIR mode now")
|
|
>>> # Example2: Usage of Static Graph
|
|
>>> import paddle
|
|
>>> import numpy as np
|
|
|
|
>>> paddle.enable_static()
|
|
|
|
>>> class MyLayer(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self.conv1 = paddle.nn.Conv2D(
|
|
... in_channels=3,
|
|
... out_channels=4,
|
|
... kernel_size=3,
|
|
... padding=2,
|
|
... )
|
|
... self.linear1 = paddle.nn.Linear(4624, 32)
|
|
... self.linear2 = paddle.nn.Linear(32, 32)
|
|
... self.linear3 = paddle.nn.Linear(32, 10)
|
|
...
|
|
... def forward(self, img):
|
|
... hidden = self.conv1(img)
|
|
... hidden = paddle.flatten(hidden, start_axis=1)
|
|
... hidden = self.linear1(hidden)
|
|
... hidden = self.linear2(hidden)
|
|
... prediction = self.linear3(hidden)
|
|
... return prediction
|
|
|
|
>>> 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, 3, 32, 32])
|
|
... label = paddle.static.data(name='label', shape=[None, 1])
|
|
... my_layer = MyLayer()
|
|
... prob = my_layer(input_data)
|
|
... loss = paddle.mean(paddle.nn.functional.square_error_cost(prob, label))
|
|
...
|
|
... optimizer = paddle.optimizer.SGD(learning_rate=0.1)
|
|
... # Calling paddle.incubate.asp.decorate() to wrap minimize() in optimizer, which
|
|
... # will insert necessary masking operations for ASP workflow.
|
|
... # In static graph mode, ASP creates related mask variables
|
|
... # during minimize().
|
|
... optimizer = paddle.incubate.asp.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 paddle.asp.prune_model()
|
|
>>> paddle.incubate.asp.prune_model(my_layer, mask_algo='mask_2d_best')
|
|
>>> # it also be accepted to call
|
|
>>> # paddle.incubate.asp.prune_model(main_program, mask_algo='mask_2d_best')
|
|
|
|
>>> for i in range(10):
|
|
... imgs = np.random.randn(64, 3, 32, 32).astype('float32')
|
|
... labels = np.random.randint(10, size=(64, 1)).astype('float32')
|
|
... exe.run(main_program, feed={'data': imgs, 'label': labels})
|
|
"""
|
|
device = paddle.device.get_device()
|
|
place = paddle.set_device(device)
|
|
|
|
MaskAlgo_mapping = {
|
|
'mask_1d': MaskAlgo.MASK_1D,
|
|
'mask_2d_greedy': MaskAlgo.MASK_2D_GREEDY,
|
|
'mask_2d_best': 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"]'
|
|
)
|
|
|
|
prune_func = None
|
|
if isinstance(model, paddle.nn.Layer):
|
|
prune_func = ASPHelper.prune_model_by_layer
|
|
elif isinstance(model, paddle.static.Program):
|
|
prune_func = ASPHelper.prune_model_by_program
|
|
if (
|
|
hasattr(model, "distributed_info_")
|
|
and model.distributed_info_["sharding_degree"] > 1
|
|
and paddle.base.is_compiled_with_cuda()
|
|
):
|
|
gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0))
|
|
place = paddle.CUDAPlace(gpu_id)
|
|
else:
|
|
raise TypeError(
|
|
f"model should be paddle.nn.Layer or paddle.static.Program, but got {type(model)}"
|
|
)
|
|
|
|
return prune_func(
|
|
place,
|
|
model,
|
|
n=n,
|
|
m=m,
|
|
mask_algo=MaskAlgo_mapping[mask_algo],
|
|
with_mask=with_mask,
|
|
)
|
|
|
|
|
|
class ProgramASPInfo:
|
|
r"""
|
|
ProgramASPInfo is a container to keep ASP relevant information of Program. It contains three inner-variables:
|
|
1. __mask_vars (Dictionary): Key is parameter's name and value is its corresponding sparse mask Variable object, which is created by `ASPHelper.create_mask_variables`.
|
|
2. __masks (Dictionary): Key is parameter's name and value is its corresponding sparse mask Numpy array, which is created by `ASPHelper.prune_model`.
|
|
3. __excluded_layers (List): It stores name of layers which should not involve into ASP workflow.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.__mask_vars = {}
|
|
self.__masks = {}
|
|
self.__excluded_layers = []
|
|
|
|
def update_mask_vars(self, param_name: str, var: Tensor) -> None:
|
|
self.__mask_vars[param_name] = var
|
|
|
|
def update_masks(self, param_name: str, var: npt.NDArray[Any]) -> None:
|
|
self.__masks[param_name] = var
|
|
|
|
def update_excluded_layers(self, param_names: list[str]) -> None:
|
|
self.__excluded_layers.extend(copy.deepcopy(param_names))
|
|
|
|
def reset_excluded_layers(self) -> None:
|
|
self.__excluded_layers = []
|
|
|
|
@property
|
|
def mask_vars(self) -> dict[str, Tensor]:
|
|
return self.__mask_vars
|
|
|
|
@property
|
|
def masks(self) -> dict[str, npt.NDArray[Any]]:
|
|
return self.__masks
|
|
|
|
@property
|
|
def excluded_layers(self) -> list[str]:
|
|
return self.__excluded_layers
|
|
|
|
|
|
class ASPHelper:
|
|
r"""
|
|
ASPHelper is a collection of Auto SParsity (ASP) functions to enable
|
|
|
|
1. training models with weights in 2:4 sparse pattern on FP16 or 1:2 sparse pattern on FP32 from scratch.
|
|
2. pruning well-trained models into 2:4 sparse pattern on FP16 or 1:2 sparse pattern on FP32 for fine-tuning.
|
|
"""
|
|
|
|
MASK_APPENDDED_NAME = 'asp_mask'
|
|
PADDLE_WEIGHT_SUFFIX = "w_"
|
|
|
|
__asp_info = {}
|
|
|
|
@classmethod
|
|
def set_excluded_layers(
|
|
cls, param_names: list[str], main_program: Program
|
|
) -> None:
|
|
r"""
|
|
This is the implementation of `asp.set_excluded_layers`, for details please see explanation in `asp.set_excluded_layers`.
|
|
"""
|
|
asp_info = cls._get_program_asp_info(main_program)
|
|
asp_info.update_excluded_layers(param_names)
|
|
|
|
@classmethod
|
|
def reset_excluded_layers(cls, main_program: Program | None = None) -> None:
|
|
r"""
|
|
This is the implementation of `asp.reset_excluded_layers`, for details please see explanation in `asp.reset_excluded_layers`.
|
|
"""
|
|
if main_program is None:
|
|
for prog in cls.__asp_info:
|
|
cls.__asp_info[prog].reset_excluded_layers()
|
|
else:
|
|
cls._get_program_asp_info(main_program).reset_excluded_layers()
|
|
|
|
@staticmethod
|
|
def decorate(optimizer: Optimizer) -> OptimizerWithSparsityGuarantee:
|
|
r"""
|
|
This is the implementation of `asp.decorate`, for details please see explanation in `asp.decorate`.
|
|
"""
|
|
if paddle.in_dynamic_mode():
|
|
# main_prog and startup_prog would be used with paddle.static.program_guard
|
|
# to create ASP masks. Moreover, main_prog is a key to map paddle.static.Program
|
|
# to its own ASP information, like ASP mask variables. For dynamic graph, we use
|
|
# default_main_program as the key.
|
|
main_prog = paddle.static.default_main_program()
|
|
startup_prog = paddle.static.default_startup_program()
|
|
ASPHelper._create_mask_variables(
|
|
main_prog, startup_prog, optimizer._parameter_list
|
|
)
|
|
return OptimizerWithSparsityGuarantee(optimizer)
|
|
|
|
@classmethod
|
|
def prune_model_by_program(
|
|
cls,
|
|
place: PlaceLike,
|
|
main_program: Program | None = None,
|
|
n: int = 2,
|
|
m: int = 4,
|
|
mask_algo: MaskAlgo = MaskAlgo.MASK_1D,
|
|
with_mask: bool = True,
|
|
) -> dict[str, npt.NDArray[Any]]:
|
|
r"""
|
|
This is the implementation of `asp.prune_model`, for details please see explanation in `asp.prune_model`.
|
|
"""
|
|
|
|
if main_program is None:
|
|
main_program = paddle.static.default_main_program()
|
|
|
|
asp_info = cls._get_program_asp_info(main_program)
|
|
for param in main_program.global_block().all_parameters():
|
|
if ASPHelper._is_supported_layer(main_program, param.name):
|
|
weight_tensor = global_scope().find_var(param.name).get_tensor()
|
|
weight_nparray = np.array(weight_tensor)
|
|
|
|
prune_func = ASPHelper._get_prune_func_by_name(param.name)
|
|
|
|
weight_pruned_nparray, weight_sparse_mask = prune_func(
|
|
weight_nparray, m, n, mask_algo, param.name
|
|
)
|
|
weight_pruned_nparray = weight_pruned_nparray.astype(
|
|
weight_nparray.dtype
|
|
)
|
|
weight_tensor.set(weight_pruned_nparray, place)
|
|
|
|
if with_mask:
|
|
weight_mask_param = global_scope().find_var(
|
|
ASPHelper._get_mask_name(param.name)
|
|
)
|
|
assert weight_mask_param is not None, (
|
|
f'Cannot find {ASPHelper._get_mask_name(param.name)} variable, please call optimizer.minimize ('
|
|
'paddle.incubate.asp.decorate(optimizer).minimize(loss)'
|
|
' and initialization (exe.run(startup_program)) first!'
|
|
)
|
|
weight_mask_tensor = weight_mask_param.get_tensor()
|
|
weight_sparse_mask = weight_sparse_mask.astype(
|
|
np.array(weight_mask_tensor).dtype
|
|
)
|
|
weight_mask_tensor.set(weight_sparse_mask, place)
|
|
asp_info.update_masks(param.name, weight_sparse_mask)
|
|
return asp_info.masks.copy()
|
|
|
|
@classmethod
|
|
def prune_model_by_layer(
|
|
cls,
|
|
place: PlaceLike,
|
|
layer: Layer,
|
|
n: int = 2,
|
|
m: int = 4,
|
|
mask_algo: MaskAlgo = MaskAlgo.MASK_1D,
|
|
with_mask: bool = True,
|
|
) -> dict[str, npt.NDArray[Any]]:
|
|
r"""
|
|
This is the implementation of `asp.prune_model`, for details please see explanation in `asp.prune_model`.
|
|
"""
|
|
if paddle.in_dynamic_mode():
|
|
main_program = paddle.static.default_main_program()
|
|
asp_info = cls._get_program_asp_info(main_program)
|
|
|
|
for param in layer.parameters():
|
|
if ASPHelper._is_supported_layer(main_program, param.name):
|
|
weight_nparray = param.numpy()
|
|
|
|
prune_func = ASPHelper._get_prune_func_by_name(param.name)
|
|
|
|
weight_pruned_nparray, weight_sparse_mask = prune_func(
|
|
weight_nparray, m, n, mask_algo, param.name
|
|
)
|
|
|
|
weight_pruned_nparray = weight_pruned_nparray.astype(
|
|
weight_nparray.dtype
|
|
)
|
|
param.set_value(weight_pruned_nparray)
|
|
|
|
if with_mask:
|
|
weight_mask_param = asp_info.mask_vars.get(
|
|
param.name, None
|
|
)
|
|
assert weight_mask_param is not None, (
|
|
f'Cannot find {ASPHelper._get_mask_name(param.name)} variable, please call asp.decorate() to'
|
|
' decorate your optimizer first!'
|
|
)
|
|
weight_mask_param.set_value(weight_sparse_mask)
|
|
|
|
asp_info.update_masks(param.name, weight_sparse_mask)
|
|
|
|
return asp_info.masks.copy()
|
|
else:
|
|
# This for loop is only used to obtain Block and Program from
|
|
# first parameters.
|
|
target_program = None
|
|
for param in layer.parameters():
|
|
target_program = param.block.program
|
|
assert target_program is not None, (
|
|
'Cannot get paddle.static.Program from Paddle.nn.Layer.'
|
|
)
|
|
return ASPHelper.prune_model_by_program(
|
|
place,
|
|
target_program,
|
|
n=n,
|
|
m=m,
|
|
mask_algo=mask_algo,
|
|
with_mask=with_mask,
|
|
)
|
|
|
|
@staticmethod
|
|
def _get_mask_name(param_name: str) -> str:
|
|
r"""
|
|
Return mask name by given parameter name :attr:`param_name`.
|
|
|
|
Args:
|
|
param_name (string): The name of parameter.
|
|
Returns:
|
|
string: The mask name of :attr:`param_name`.
|
|
"""
|
|
return param_name + "." + ASPHelper.MASK_APPENDDED_NAME
|
|
|
|
@staticmethod
|
|
def _get_not_ASP_relevant_vars(main_program: Program) -> list[Tensor]:
|
|
r"""
|
|
Get all parameters's Variables in :attr:`main_program` but excluded ASP mask Variables.
|
|
|
|
Args:
|
|
main_program (Program): Program with model definition and its parameters.
|
|
Returns:
|
|
list: A list of parameter Variables in :attr:`main_program` (excluded ASP mask Variables).
|
|
"""
|
|
var_list = []
|
|
for param in main_program.global_block().all_parameters():
|
|
param_name_list = param.name.split('.')
|
|
|
|
if ASPHelper.MASK_APPENDDED_NAME not in param_name_list:
|
|
var_list.append(param)
|
|
return var_list
|
|
|
|
@classmethod
|
|
def _get_program_asp_info(cls, main_program: Program) -> ProgramASPInfo:
|
|
if main_program not in cls.__asp_info:
|
|
cls.__asp_info[main_program] = ProgramASPInfo()
|
|
return cls.__asp_info[main_program]
|
|
|
|
@classmethod
|
|
def _is_supported_layer(
|
|
cls, main_program: Program, param_name: str
|
|
) -> bool:
|
|
r"""
|
|
Verify if given :attr:`param_name` is supported by ASP.
|
|
|
|
Args:
|
|
param_name (string): The name of parameter.
|
|
Returns:
|
|
bool: True if it is supported, else False.
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +SKIP("paddle.incubate.asp module doesn't support PIR mode now")
|
|
>>> from paddle.incubate.asp import ASPHelper
|
|
>>> 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])
|
|
... fc = paddle.static.nn.fc(x=input_data, num_flatten_dims=-1, size=32, activation=None)
|
|
|
|
>>> for param in main_program.global_block().all_parameters():
|
|
... print(param.name, '->', ASPHelper._is_supported_layer(main_program, param.name))
|
|
fc_0.w_0 -> True
|
|
fc_0.b_0 -> False
|
|
"""
|
|
param_name_list = param_name.split('.')
|
|
|
|
if ASPHelper.MASK_APPENDDED_NAME in param_name_list:
|
|
return False
|
|
|
|
for layer in cls._get_program_asp_info(main_program).excluded_layers:
|
|
if layer in param_name:
|
|
return False
|
|
|
|
if param_name in supported_layers_and_prune_func_map:
|
|
return True
|
|
|
|
# The parameter's name is neither in *.* format nor added to supported_layers_and_prune_func_map, return False.
|
|
if len(param_name_list) == 1:
|
|
return False
|
|
|
|
param_name_no_weight_suffix = param_name_list[0]
|
|
param_type_suffix = param_name_list[1]
|
|
layer_name = param_name_no_weight_suffix[
|
|
: param_name_no_weight_suffix.rfind('_')
|
|
]
|
|
if ASPHelper.PADDLE_WEIGHT_SUFFIX not in param_type_suffix:
|
|
return False
|
|
|
|
if (
|
|
param_name_no_weight_suffix in supported_layers_and_prune_func_map
|
|
or layer_name in supported_layers_and_prune_func_map
|
|
):
|
|
return True
|
|
|
|
return False
|
|
|
|
@classmethod
|
|
def _get_prune_func_by_name(
|
|
cls, param_name: str
|
|
) -> Callable[
|
|
[npt.NDArray[Any], int, int, MaskAlgo, str],
|
|
tuple[npt.NDArray[Any], npt.NDArray[Any]],
|
|
]:
|
|
func = supported_layers_and_prune_func_map.get(param_name, None)
|
|
param_name_no_weight_suffix = param_name.split('.')[0]
|
|
if func is None:
|
|
func = supported_layers_and_prune_func_map.get(
|
|
param_name_no_weight_suffix, None
|
|
)
|
|
if func is None:
|
|
layer_name = param_name_no_weight_suffix[
|
|
: param_name_no_weight_suffix.rfind('_')
|
|
]
|
|
func = supported_layers_and_prune_func_map.get(
|
|
layer_name, _default_pruning
|
|
)
|
|
return func
|
|
|
|
@classmethod
|
|
def _minimize(
|
|
cls,
|
|
optimizer: Optimizer,
|
|
loss: Tensor,
|
|
main_program: Program | None = None,
|
|
startup_program: Program | None = None,
|
|
parameter_list: Iterable[Tensor] | Iterable[str] | None = None,
|
|
no_grad_set: set[Tensor] | set[str] | None = None,
|
|
) -> tuple[list[Operator], list[tuple[Tensor, Tensor]]]:
|
|
r"""
|
|
This function is a decorator of `minimize` function in `Optimizer`.
|
|
There are three steps:
|
|
|
|
1. Call :attr:`optimizer`.minimize(:attr:`loss`)
|
|
2. Create sparse mask Tensors according to supported layers in :attr:`main_program`.
|
|
3. Insert masking ops in the end of parameters update.
|
|
|
|
*Note*: Please use `ASP.decorate` instead when applying distributed training with `Fleet`.
|
|
(Due to there is a invisible graphs optimization in `Fleet.minimize()` which make training graph
|
|
cannot be modified anymore.)
|
|
|
|
Args:
|
|
optimizer (Optimizer): A Optimizer used for training.
|
|
loss (Variable): A Variable containing the value to minimize.
|
|
main_program (Program, optional): Program with model definition and its parameters. Default is `loss.block.program`.
|
|
startup_program (Program, optional): Program for initializing parameters in `parameter_list`. Default is `paddle.static.default_startup_program()`.
|
|
parameter_list (Iterable, optional): Iterable of `Variable` or `Variable.name` to update to minimize `loss`. The default value is None, at this time all parameters will be updated.
|
|
no_grad_set (set, optional): Set of `Variable or `Variable.name` that don't need to be updated. The default value is None.
|
|
Returns:
|
|
list: operators from :attr:`optimizer`.minimize(:attr:`loss`).
|
|
list: pairs of parameters and their gradients.
|
|
"""
|
|
if main_program is None:
|
|
main_program = loss.block.program
|
|
|
|
if startup_program is None:
|
|
startup_program = paddle.static.default_startup_program()
|
|
|
|
optimizer_ops, params_and_grads = optimizer.minimize(
|
|
loss, startup_program, parameter_list, no_grad_set=no_grad_set
|
|
)
|
|
|
|
params_only = [pg[0] for pg in params_and_grads]
|
|
cls._create_mask_variables(main_program, startup_program, params_only)
|
|
cls._insert_sparse_mask_ops(main_program, params_only)
|
|
return optimizer_ops, params_and_grads
|
|
|
|
@classmethod
|
|
@dygraph_only
|
|
def _step(cls, optimizer: Optimizer) -> None:
|
|
r"""
|
|
This function is a decorator of `step` function in `Optimizer`.
|
|
There are three steps:
|
|
|
|
1. Call :attr:`optimizer`.step()
|
|
2. Mask parameters with sparse masks.
|
|
|
|
*Note*: Please use `ASP.decorate` instead when applying distributed training with `Fleet`.
|
|
(Due to there is a invisible graphs optimization in `Fleet.minimize()` which make training graph
|
|
cannot be modified anymore.)
|
|
|
|
Args:
|
|
optimizer (Optimizer): A Optimizer used for training.
|
|
"""
|
|
optimizer.step()
|
|
main_prog = paddle.static.default_main_program()
|
|
with paddle.base.dygraph.no_grad():
|
|
ASPHelper._insert_sparse_mask_ops(
|
|
main_prog, optimizer._parameter_list
|
|
)
|
|
|
|
@classmethod
|
|
def _create_mask_variables(
|
|
cls,
|
|
main_program: Program,
|
|
startup_program: Program,
|
|
params: Sequence[Tensor],
|
|
) -> None:
|
|
r"""
|
|
Create sparse mask Tensors according to supported layers in :attr:`main_program`.
|
|
This function is called in second step of `ASPHelper._minimize`
|
|
|
|
Args:
|
|
main_program (Program): Program with model definition and its parameters.
|
|
startup_program (Program): Program for initializing parameters.
|
|
params (list): Variable parameters.
|
|
"""
|
|
asp_info = cls._get_program_asp_info(main_program)
|
|
with program_guard(main_program, startup_program):
|
|
for param in params:
|
|
if ASPHelper._is_supported_layer(main_program, param.name):
|
|
if param.name not in asp_info.mask_vars:
|
|
mask_param = paddle.create_parameter(
|
|
name=ASPHelper._get_mask_name(param.name),
|
|
shape=param.shape,
|
|
dtype=param.dtype,
|
|
default_initializer=paddle.nn.initializer.Constant(
|
|
value=1.0
|
|
),
|
|
)
|
|
mask_param.stop_gradient = True
|
|
mask_param.trainable = False
|
|
asp_info.update_mask_vars(param.name, mask_param)
|
|
|
|
@classmethod
|
|
def _insert_sparse_mask_ops(
|
|
cls, main_program: Program, params: Sequence[Tensor]
|
|
) -> None:
|
|
r"""
|
|
Insert masking ops in the end of parameters update.
|
|
This function is called in third step of `ASPHelper._minimize`
|
|
|
|
Args:
|
|
main_program (Program): Program with model definition and its parameters.
|
|
params (list): Variable parameters.
|
|
"""
|
|
block = main_program.global_block()
|
|
asp_info = cls._get_program_asp_info(main_program)
|
|
for param in params:
|
|
if param.name in asp_info.mask_vars:
|
|
block.append_op(
|
|
type='elementwise_mul',
|
|
inputs={"X": param, 'Y': asp_info.mask_vars[param.name]},
|
|
outputs={'Out': param},
|
|
attrs={
|
|
'axis': -1,
|
|
OP_ROLE_KEY: int(OpRole.Optimize),
|
|
},
|
|
)
|
|
|
|
|
|
class OptimizerWithSparsityGuarantee:
|
|
r"""
|
|
OptimizerWithSparsityGuarantee is a wrapper to decorate `minimize` function of given optimizer by `_minimize` of ASPHelper.
|
|
The decorated `minimize` function would do three things (exactly same as `ASPHelper._minimize`):
|
|
1. Call `minimize` function of given optimizer.
|
|
2. Call `ASPHelper._create_mask_variables` to create mask Variables.
|
|
3. Call `ASPHelper._insert_sparse_mask_ops` to insert weight masking ops in the end of `loss`'s Program.
|
|
"""
|
|
|
|
def __init__(self, optimizer: Optimizer) -> None:
|
|
self._optimizer = optimizer
|
|
|
|
def __getattr__(self, item: str) -> Any:
|
|
return getattr(self._optimizer, item)
|
|
|
|
def minimize(
|
|
self,
|
|
loss: Tensor,
|
|
startup_program: Program | None = None,
|
|
parameter_list: Iterable[Tensor] | Iterable[str] | None = None,
|
|
no_grad_set: set[Tensor] | set[str] | None = None,
|
|
) -> tuple[list[Operator], list[tuple[Tensor, Tensor]]]:
|
|
r"""
|
|
This function is to call `ASPHelper.minimize()` and return its return
|
|
|
|
Args:
|
|
loss (Variable): A Variable containing the value to minimize.
|
|
startup_program (Program, optional): Program for initializing parameters in `parameter_list`. Default is `paddle.static.default_startup_program()`.
|
|
parameter_list (Iterable, optional): Iterable of `Variable` or `Variable.name` to update to minimize `loss`. The default value is None, at this time all parameters will be updated.
|
|
no_grad_set (set, optional): Set of `Variable or `Variable.name` that don't need to be updated. The default value is None.
|
|
Returns:
|
|
list: operators from :attr:`optimizer`.minimize(:attr:`loss`).
|
|
list: pairs of parameters and their gradients.
|
|
"""
|
|
return ASPHelper._minimize(
|
|
self._optimizer,
|
|
loss,
|
|
startup_program=startup_program,
|
|
parameter_list=parameter_list,
|
|
no_grad_set=no_grad_set,
|
|
)
|
|
|
|
@dygraph_only
|
|
def step(self) -> None:
|
|
r"""
|
|
This function is a decorator of `step` function in `Optimizer`.
|
|
There are three steps:
|
|
|
|
1. Call :attr:`optimizer`.step()
|
|
2. Mask parameters with sparse masks.
|
|
|
|
*Note*: Please use `ASP.decorate` instead when applying distributed training with `Fleet`.
|
|
(Due to there is a invisible graphs optimization in `Fleet.minimize()` which make training graph
|
|
cannot be modified anymore.)
|
|
|
|
Args:
|
|
optimizer (Optimizer): A Optimizer used for training.
|
|
"""
|
|
ASPHelper._step(self._optimizer)
|
|
|
|
@dygraph_only
|
|
def state_dict(self) -> dict[str, Tensor]:
|
|
r"""
|
|
This function is a decorator of `state_dict` function in `Optimizer`.
|
|
|
|
Returns:
|
|
state_dict(dict) : dict contains all the Tensor used by optimizer
|
|
"""
|
|
state_dict = self._optimizer.state_dict()
|
|
asp_info = ASPHelper._get_program_asp_info(
|
|
paddle.static.default_main_program()
|
|
)
|
|
for param_name, var in asp_info.mask_vars.items():
|
|
state_dict.update({ASPHelper._get_mask_name(param_name): var})
|
|
return state_dict
|
|
|
|
@dygraph_only
|
|
def set_state_dict(self, state_dict: dict[str, Tensor]) -> None:
|
|
r"""
|
|
This function is a decorator of `set_state_dict` function in `Optimizer`.
|
|
Args:
|
|
state_dict(dict) : Dict contains all the Tensor needed by optimizer
|
|
Return:
|
|
None
|
|
"""
|
|
asp_info = ASPHelper._get_program_asp_info(
|
|
paddle.static.default_main_program()
|
|
)
|
|
for param_name, var in asp_info.mask_vars.items():
|
|
param_mask_name = ASPHelper._get_mask_name(param_name)
|
|
assert param_mask_name in state_dict, (
|
|
f"The {param_mask_name} is not found."
|
|
)
|
|
var.set_value(state_dict[param_mask_name])
|
|
asp_info.update_masks(param_name, var.numpy())
|
|
return self._optimizer.set_state_dict(state_dict)
|