Spaces:
Sleeping
Sleeping
| # Copyright (c) ONNX Project Contributors | |
| # SPDX-License-Identifier: Apache-2.0 | |
| from typing import Optional, Tuple | |
| import numpy as np | |
| from numpy.random import RandomState # type: ignore | |
| from onnx.reference.op_run import OpRun | |
| def _dropout( | |
| X: np.ndarray, | |
| drop_probability: float = 0.5, | |
| seed: Optional[int] = None, | |
| training_mode: bool = False, | |
| return_mask: bool = False, | |
| ) -> Tuple[np.ndarray]: | |
| if drop_probability == 0 or not training_mode: | |
| if return_mask: | |
| return X, np.ones(X.shape, dtype=bool) # type: ignore | |
| return (X,) | |
| rnd = RandomState(seed) | |
| mask = rnd.uniform(0, 1.0, X.shape) >= drop_probability | |
| scale = 1.0 / (1.0 - drop_probability) | |
| return (mask * X * scale, mask.astype(bool)) if return_mask else (mask * X * scale,) # type: ignore | |
| class DropoutBase(OpRun): | |
| def __init__(self, onnx_node, run_params): # type: ignore | |
| OpRun.__init__(self, onnx_node, run_params) | |
| self.n_outputs = len(onnx_node.output) | |
| def _private_run( | |
| self, | |
| X: np.ndarray, | |
| seed: Optional[int] = None, | |
| ratio: float = 0.5, | |
| training_mode: bool = False, | |
| ) -> Tuple[np.ndarray]: | |
| return _dropout( | |
| X, | |
| ratio, | |
| seed=seed, # type: ignore | |
| return_mask=self.n_outputs == 2, | |
| training_mode=training_mode, | |
| ) | |
| class Dropout_7(DropoutBase): | |
| def _run(self, X, ratio=None): # type: ignore | |
| return self._private_run(X, ratio) | |
| class Dropout_12(DropoutBase): | |
| def _run(self, *inputs, seed=None): # type: ignore | |
| X = inputs[0] | |
| ratio = 0.5 if len(inputs) <= 1 else inputs[1] | |
| training_mode = False if len(inputs) <= 2 else inputs[2] | |
| return self._private_run( | |
| X, seed=seed, ratio=ratio, training_mode=training_mode # type: ignore | |
| ) | |