helper.py 11.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
# Copyright (c) 2022 PaddlePaddle Authors. 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.

import logging
16
from collections import defaultdict
17

18 19
import paddle

20 21 22 23
from paddle.nn import Layer
from paddle.jit import to_static, not_to_static
from paddle.distributed.utils import get_logger
from paddle.fluid.framework import Operator, Parameter, _non_static_mode
24
from paddle.fluid.framework import program_guard
25
from paddle.fluid.executor import global_scope
26
from paddle.fluid.dygraph.dygraph_to_static.program_translator import StaticFunction
27 28

from .utils import to_list
29
from .converter import Converter
30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48


class ProxyLayer(Layer):
    """
    ProxyLayer implements all logic for converting dygraph model into
    static Program IR. Meanwhile, it provides conviential interfaces for
    auto parallel to visit feed/fetch/loss/metric variables.
    """

    def __init__(self, layer, loss_func, metrics):
        super(ProxyLayer, self).__init__()
        # NOTE: All verify logics are finished in Engine.Prepare
        self.inner_layer = layer
        self.loss_func = loss_func
        self.metrics = metrics
        # train / eval / predict
        self.mode = None

        # generated program vars
49 50 51 52 53
        self._input_vars = defaultdict(list)
        self._label_vars = defaultdict(list)
        self._output_vars = defaultdict(list)
        self._loss_vars = defaultdict(list)
        self._metric_vars = defaultdict(list)
54 55 56 57 58 59

    def _train(self, inputs, labels):
        """
        Train process of inner_layer with forward/loss/metric logic.
        """
        # step 1. save feed variables of Program
60 61 62
        mode = 'train'
        self._input_vars[mode] = inputs
        self._label_vars[mode] = labels
63 64

        # step 2. call inner_layer.forward
65
        self._output_vars[mode] = self.inner_layer(*inputs)
66 67 68

        # step 3. calculate loss if needed
        new_inputs = self._prepare(self.output_vars, labels)
69
        self._loss_vars[mode] = self.call_loss(new_inputs)
70 71

        # step 4. calculate metrics if needed
72
        self._metric_vars[mode] = self.call_metrics(new_inputs)
73 74 75 76 77 78 79 80 81

    def _eval(self, inputs, labels):
        """
        Evaluate process of inner_layer with forward/loss/metric logic.
        """
        # TODO(dev): we can reuse codes with self._train after making
        # sure if they can.

        # step 1. save feed variables of Program
82 83 84
        mode = 'eval'
        self._input_vars[mode] = inputs
        self._label_vars[mode] = labels
85 86

        # step 2. call inner_layer.forward
87
        self._output_vars[mode] = self.inner_layer(*inputs)
88 89 90

        # step 3. calculate loss if needed
        new_inputs = self._prepare(self.output_vars, labels)
91
        self._loss_vars[mode] = self.call_loss(new_inputs)
92 93

        # step 4. calculate metrics if needed
94
        self._metric_vars[mode] = self.call_metrics(new_inputs)
95

96
    def _predict(self, inputs, labels):
97 98 99 100
        """
        Predict process of inner_layer with forward logic.
        """
        # step 1. save feed variables of Program
101 102
        mode = 'predict'
        self._input_vars[mode] = inputs
103
        self._label_vars[mode] = labels
104 105

        # step 2. call inner_layer.forward
106
        self._output_vars[mode] = self.inner_layer(*inputs)
107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149

    @not_to_static
    def _prepare(self, outputs, labels):
        """
        Concat outputs and labels as a single list

        NOTE(dev): We use @not_to_static to avoid AST Analysis.
        """
        return to_list(outputs) + to_list(labels)

    def call_loss(self, inputs):
        """
        Apply Loss Function on outputs and labels.

        Args:
            inputs: List[Variable]

        Returns: List[Variable]
        """
        res = []
        if self.loss_func is not None:
            res = self.loss_func(*inputs)
        return res

    def call_metrics(self, inputs):
        """
        Apply Metrics Function on outputs and labels.

        Args:
            inputs: List[Variable]

        Returns: List[Variable]
        """
        outs = []
        for metric in self.metrics:
            outs.extend(metric.compute(*inputs))

        return outs

    def set_mode(self, mode):
        self.mode = mode
        self.training = mode == 'train'

150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172
    def clone(self):
        return ProxyLayer(self.inner_layer, self.loss_func, self.metrics)

    @property
    def input_vars(self):
        return self._input_vars[self.mode]

    @property
    def label_vars(self):
        return self._label_vars[self.mode]

    @property
    def output_vars(self):
        return self._output_vars[self.mode]

    @property
    def loss_vars(self):
        return self._loss_vars[self.mode]

    @property
    def metric_vars(self):
        return self._metric_vars[self.mode]

173 174 175 176
    @property
    def startup_program(self):
        return self.inner_layer._startup_program()

177 178 179

class BuildInfo:

180 181 182 183 184 185 186 187
    def __init__(self):
        self.clear()

    def has_cache(self, mode, update=False):
        is_cache = self.states[mode]
        if update:
            self.cache(mode)
        return is_cache
188

189 190 191 192 193
    def cache(self, mode):
        self.states[mode] = True

    def clear(self):
        self.states = defaultdict(bool)
194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210


class ProgramHelper(object):
    """
    A Helper class for Engine to provides different Program IR according specified 'mode'.
    """

    def __init__(self, layer, loss_func, metrics, inputs_spec, labels_spec):
        # original model config information
        # TODO(Aurelius84): Implenet append_backward and optimizer in ProxyLayer
        # after distribute engine satisify basic condition.
        self.proxy_layer = ProxyLayer(layer, loss_func, metrics)
        self.inputs_spec = inputs_spec
        self.labels_spec = labels_spec

        self.build_info = BuildInfo()
        self._logger = get_logger(logging.INFO)
211
        self.lazy_init = False
212

213 214 215 216 217 218 219
    def reset(self):
        """
        Reset all state of current Object.
        """
        self.build_info.clear()
        self.proxy_layer = self.proxy_layer.clone()

220 221 222 223 224
    def build_program(self, mode):
        """
        Convert dygraph model into static Program IR.
        """
        assert mode in ['train', 'eval', 'predict']
225
        self.proxy_layer.set_mode(mode)
226
        # skip if we has already built program.
227
        if self.build_info.has_cache(mode, True):
228 229 230 231 232 233
            self._logger.info(
                "Already build program with mode = %s, use cached program." %
                mode)
            return

        self._logger.info("start to build program for mode = %s." % mode)
234
        input_spec = [self.inputs_spec, self.labels_spec]
235 236 237 238 239 240 241 242 243
        static_func = to_static(self.static_func(), input_spec=input_spec)

        func_name = '_' + mode
        setattr(self.proxy_layer, func_name, static_func)

        # NOTE(dev): Because @to_static is a Lazy mechanism, so we explicitly call this to trigger
        # generating Program IR immediately.
        getattr(self.proxy_layer, func_name).concrete_program

244 245
        self._build_startup_program()

246 247 248 249
    def _build_startup_program(self):
        """
        Create and Sync parameters into startup program.
        """
250 251 252
        if len(self.startup_program.global_block().ops) > 1:
            self.lazy_init = True
            return
253 254 255 256 257 258 259 260 261
        for param in self.concrete_program.parameters:
            Parameter(name=param.name,
                      desc=param,
                      type=param.type,
                      shape=param.shape,
                      dtype=param.dtype,
                      stop_gradient=param.stop_gradient,
                      block=self.startup_program.global_block())

262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298
    def apply_optimizer(self, optimizer):
        """
        Append backward and generate optimizer operations.
        """
        self._verify_optimizer(optimizer)
        self._logger.info("start to apply optimizer: %s ",
                          type(optimizer).__name__)
        # clear optimizer parameters
        original_params = optimizer._parameter_list
        optimizer._parameter_list = None
        with program_guard(self.main_program, self.startup_program):
            res = optimizer.minimize(self.loss_vars[0])

        # restore optimizer parameters
        optimizer._parameter_list = original_params
        return res

    def _verify_optimizer(self, optimizer):
        assert optimizer is not None
        assert hasattr(optimizer,
                       "minimize"), "Optimizer must have minimize() method."
        assert self.proxy_layer.mode == 'train', "Required mode == 'train', but received '%s'" % self.proxy_layer.mode
        assert len(
            self.loss_vars
        ) == 1, "Required len(loss_vars) == 1, but received len(loss_vars) = %s" % len(
            self.loss_vars)

    def to(self, mode):
        """
        Switch underly proxy layer mode into target mode.
        """
        assert mode in ['train', 'eval', 'predict']
        func = getattr(self.proxy_layer, '_' + mode)
        assert isinstance(
            func, StaticFunction), "Please call build_program(mode) firstly."
        self.proxy_layer.set_mode(mode)

299 300
    def static_func(self):
        """
301
        Return StaticFunction instance with underly target mode.
302 303 304 305 306 307 308
        """
        assert self.proxy_layer.mode in [
            'train', 'eval', 'predict'
        ], "Please call build_program(mode) firstly."
        func_name = '_' + self.proxy_layer.mode
        return getattr(self.proxy_layer, func_name)

309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330
    def init(self, main_program, place, dist_context):
        if self.lazy_init:
            return
        for param in self.concrete_program.parameters:
            # create var in scope and share parameters to scope
            if param.name not in main_program.global_block().vars:
                continue
            # get param_var's dist_attr
            var = main_program.global_block().vars[param.name]
            var_dist_attr = dist_context.get_tensor_dist_attr_for_program(var)
            dist_attr = {
                "dims_mapping": var_dist_attr.dims_mapping,
                "process_shape": var_dist_attr.process_mesh.topology,
                "process_group": var_dist_attr.process_mesh.processes
            }
            # slice param_value with dist_attr
            # share sliced_param_value with param_tensor in global_scope
            param_tensor = global_scope().var(param.name).get_tensor()
            sliced_param = Converter.slice_with_dist_attr(
                param.numpy(), dist_attr)
            param_tensor.set(sliced_param, place)

331 332 333 334 335 336 337 338 339 340
    @property
    def concrete_program(self):
        return self.static_func().concrete_program

    @property
    def main_program(self):
        return self.concrete_program.main_program

    @property
    def startup_program(self):
341 342 343 344 345 346
        try:
            return self.proxy_layer.startup_program
        except Exception as err:
            if isinstance(err, AssertionError):
                return self.concrete_program.startup_program
            raise err
347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366

    @property
    def input_vars(self):
        return to_list(self.proxy_layer.input_vars)

    @property
    def output_vars(self):
        return to_list(self.proxy_layer.output_vars)

    @property
    def label_vars(self):
        return to_list(self.proxy_layer.label_vars)

    @property
    def loss_vars(self):
        return to_list(self.proxy_layer.loss_vars)

    @property
    def metric_vars(self):
        return to_list(self.proxy_layer.metric_vars)