wehub-resource-sync
312 lines
10 KiB
Python
312 lines
10 KiB
Python
# Copyright (c) 2023 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.
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, overload
|
|
|
|
import paddle
|
|
from paddle import base
|
|
from paddle.base import core
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Sequence
|
|
from typing import Literal, Protocol
|
|
|
|
from paddle._typing import PlaceLike
|
|
|
|
class _GeneratorState(Protocol):
|
|
def current_seed(self) -> int: ...
|
|
|
|
|
|
__all__ = []
|
|
|
|
|
|
@overload
|
|
def get_rng_state(
|
|
device: PlaceLike | None = ..., use_index: Literal[True] = ...
|
|
) -> list[int]: ...
|
|
|
|
|
|
@overload
|
|
def get_rng_state(
|
|
device: PlaceLike | None = ..., use_index: Literal[False] = ...
|
|
) -> list[_GeneratorState]: ...
|
|
|
|
|
|
@overload
|
|
def get_rng_state(
|
|
device: PlaceLike | None = ..., use_index: bool = ...
|
|
) -> list[int] | list[_GeneratorState]: ...
|
|
|
|
|
|
def get_rng_state(device=None, use_index=False):
|
|
"""
|
|
Get all random states of random generators of specified device.
|
|
|
|
Args:
|
|
device(str): This parameter determines the specific running device.
|
|
It can be ``cpu``, ``gpu``, ``xpu``, Default is None.
|
|
If None, return the generators of current device (specified by ``set_device``).
|
|
use_index(bool): If use index is True, return the index that saved in the generator
|
|
|
|
Returns:
|
|
GeneratorState: object.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> sts = paddle.incubate.get_rng_state()
|
|
"""
|
|
|
|
def get_state(generator):
|
|
if use_index:
|
|
return generator.get_state_index()
|
|
else:
|
|
return generator.get_state()
|
|
|
|
state_list = []
|
|
if device is None:
|
|
place = base.framework._current_expected_place_()
|
|
else:
|
|
place = paddle.device._convert_to_place(device)
|
|
|
|
if isinstance(place, core.CPUPlace):
|
|
state_list.append(get_state(core.default_cpu_generator()))
|
|
elif isinstance(place, core.CUDAPlace):
|
|
for i in range(core.get_cuda_device_count()):
|
|
state_list.append(get_state(core.default_cuda_generator(i)))
|
|
elif isinstance(place, core.XPUPlace):
|
|
for i in range(core.get_xpu_device_count()):
|
|
state_list.append(get_state(core.default_xpu_generator(i)))
|
|
elif isinstance(place, core.CustomPlace):
|
|
dev_cnt = sum(
|
|
[
|
|
place.get_device_type() == s.split(':')[0]
|
|
for s in core.get_available_custom_device()
|
|
]
|
|
)
|
|
for i in range(dev_cnt):
|
|
state_list.append(
|
|
get_state(
|
|
core.default_custom_device_generator(
|
|
core.CustomPlace(place.get_device_type(), i)
|
|
)
|
|
)
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"get_rng_state is not implemented for current device: {place}"
|
|
)
|
|
|
|
return state_list
|
|
|
|
|
|
@overload
|
|
def set_rng_state(
|
|
state_list: Sequence[int],
|
|
device: PlaceLike | None = ...,
|
|
use_index: Literal[True] = ...,
|
|
) -> None: ...
|
|
|
|
|
|
@overload
|
|
def set_rng_state(
|
|
state_list: Sequence[_GeneratorState],
|
|
device: PlaceLike | None = ...,
|
|
use_index: Literal[False] = ...,
|
|
) -> None: ...
|
|
|
|
|
|
@overload
|
|
def set_rng_state(
|
|
state_list: Sequence[int] | Sequence[_GeneratorState],
|
|
device: PlaceLike | None = ...,
|
|
use_index: bool = ...,
|
|
) -> None: ...
|
|
|
|
|
|
def set_rng_state(
|
|
state_list,
|
|
device=None,
|
|
use_index=False,
|
|
):
|
|
"""
|
|
|
|
Sets generator state for all device generators.
|
|
|
|
Args:
|
|
state_list(list|tuple): The device states to set back to device generators. state_list is obtained from get_rng_state().
|
|
device(str): This parameter determines the specific running device.
|
|
It can be ``cpu``, ``gpu``, ``xpu``, Default is None.
|
|
If None, return the generators of current device (specified by ``set_device``).
|
|
use_index(bool): If use index is True, state_list should be the indices of the states
|
|
|
|
Returns:
|
|
None.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> sts = paddle.incubate.get_rng_state()
|
|
>>> paddle.incubate.set_rng_state(sts)
|
|
|
|
"""
|
|
|
|
def set_state(generator, state):
|
|
if use_index:
|
|
generator.set_state_index(state)
|
|
else:
|
|
generator.set_state(state)
|
|
|
|
if device is None:
|
|
place = base.framework._current_expected_place_()
|
|
else:
|
|
place = device._convert_to_place(device)
|
|
|
|
if isinstance(place, core.CUDAPlace):
|
|
if not len(state_list) == core.get_cuda_device_count():
|
|
raise ValueError(
|
|
"Length of gpu state list should be equal to the gpu device count"
|
|
)
|
|
for i in range(core.get_cuda_device_count()):
|
|
set_state(core.default_cuda_generator(i), state_list[i])
|
|
elif isinstance(place, core.XPUPlace):
|
|
if not len(state_list) == core.get_xpu_device_count():
|
|
raise ValueError(
|
|
"Length of xpu state list should be equal to the xpu device count"
|
|
)
|
|
for i in range(core.get_xpu_device_count()):
|
|
set_state(core.default_xpu_generator(i), state_list[i])
|
|
elif isinstance(place, core.CustomPlace):
|
|
dev_types = core.get_all_custom_device_type()
|
|
dev_type = dev_types[0]
|
|
dev_cnt = core.get_custom_device_count(dev_type)
|
|
if not len(state_list) == dev_cnt:
|
|
raise ValueError(
|
|
f"Length of custom device state list should be equal to the {dev_cnt} device count"
|
|
)
|
|
for i in range(dev_cnt):
|
|
set_state(
|
|
core.default_custom_device_generator(
|
|
core.CustomPlace(place.get_device_type(), i)
|
|
),
|
|
state_list[i],
|
|
)
|
|
elif isinstance(place, core.CPUPlace):
|
|
if not len(state_list) == 1:
|
|
raise ValueError("Length of cpu state list should be equal to 1")
|
|
set_state(core.default_cpu_generator(), state_list[0])
|
|
else:
|
|
raise ValueError(
|
|
f"set_rng_state is not implemented for current device: {place}"
|
|
)
|
|
|
|
|
|
def register_rng_state_as_index(
|
|
state_list: Sequence[_GeneratorState] | None = None,
|
|
device: PlaceLike | None = None,
|
|
) -> list[int]:
|
|
"""
|
|
|
|
The register_rng_state_as_index function creates and registers a new generator state within the generator.
|
|
It enables users to manage multiple generator states via indices,
|
|
offering a convenient way to switch between these states without directly manipulating the generator's state.
|
|
|
|
Args:
|
|
state_list(list|tuple): A list or tuple representing the RNG states for devices.
|
|
If not provided, the function will register the current state.
|
|
device(str): This parameter determines the specific running device.
|
|
It can be ``cpu``, ``gpu``, ``xpu``, Default is None.
|
|
If None, return the generators of current device (specified by ``set_device``).
|
|
|
|
Returns:
|
|
A list of indices representing the positions at which the new states were saved within the generator.
|
|
These indices can be used to switch between states using set_rng_state(use_index=True)
|
|
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> old_index = paddle.incubate.get_rng_state(use_index=True)
|
|
>>> print(old_index)
|
|
[0]
|
|
>>> new_index = paddle.incubate.register_rng_state_as_index()
|
|
>>> print(new_index)
|
|
[1]
|
|
>>> paddle.incubate.set_rng_state(old_index, use_index=True)
|
|
>>> paddle.incubate.set_rng_state(new_index, use_index=True)
|
|
|
|
"""
|
|
new_state_index_list = []
|
|
|
|
if device is None:
|
|
place = base.framework._current_expected_place()
|
|
else:
|
|
place = device._convert_to_place(device)
|
|
|
|
if state_list is None:
|
|
state_list = get_rng_state(device)
|
|
|
|
if isinstance(place, core.CUDAPlace):
|
|
if not len(state_list) == core.get_cuda_device_count():
|
|
raise ValueError(
|
|
"Length of gpu state list should be equal to the gpu device count"
|
|
)
|
|
for i in range(core.get_cuda_device_count()):
|
|
new_state_index_list.append(
|
|
core.default_cuda_generator(i).register_state_index(
|
|
state_list[i]
|
|
)
|
|
)
|
|
elif isinstance(place, core.XPUPlace):
|
|
if not len(state_list) == core.get_xpu_device_count():
|
|
raise ValueError(
|
|
"Length of xpu state list should be equal to the xpu device count"
|
|
)
|
|
for i in range(core.get_xpu_device_count()):
|
|
new_state_index_list.append(
|
|
core.default_xpu_generator(i).register_state_index(
|
|
state_list[i]
|
|
)
|
|
)
|
|
elif isinstance(place, core.CustomPlace):
|
|
dev_types = core.get_all_custom_device_type()
|
|
dev_type = dev_types[0]
|
|
dev_cnt = core.get_custom_device_count(dev_type)
|
|
if not len(state_list) == dev_cnt:
|
|
raise ValueError(
|
|
f"Length of custom device state list should be equal to the {dev_cnt} device count"
|
|
)
|
|
for i in range(dev_cnt):
|
|
new_state_index_list.append(
|
|
core.default_custom_device_generator(
|
|
core.CustomPlace(place.get_device_type(), i)
|
|
).register_state_index(state_list[i])
|
|
)
|
|
elif isinstance(place, core.CPUPlace):
|
|
if not len(state_list) == 1:
|
|
raise ValueError("Length of cpu state list should be equal to 1")
|
|
new_state_index_list.append(
|
|
core.default_cpu_generator().register_state_index(state_list[0])
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"register_rng_state_index is not implemented for current device: {place}"
|
|
)
|
|
return new_state_index_list
|