mirror of
https://github.com/wassname/pyrobolearn.git
synced 2026-09-09 11:31:38 +08:00
update storages + ER
This commit is contained in:
Brian Delhaisse
@@ -14,7 +14,7 @@ The framework has been tested with Python 2.7 and Ubuntu 16.04 and 18.04. We als
|
||||
## Installation
|
||||
|
||||
1. First download the `pip` Python package manager and create a virtual environment for Python as described in the following link: https://packaging.python.org/guides/installing-using-pip-and-virtualenv/
|
||||
On Ubuntu, in the terminal, you can type to download and install `pip` and `virtualenv`:
|
||||
On Ubuntu, you can install `pip` and `virtualenv` by typing in the terminal:
|
||||
|
||||
- In Python 2.7:
|
||||
```bash
|
||||
@@ -34,14 +34,14 @@ virtualenv -p /usr/bin/python<version> <virtualenv_name>
|
||||
# activate the virtual environment
|
||||
source <virtualenv_name>/bin/activate
|
||||
```
|
||||
where `<version>` is the python version you want to use (select between `2.7` or `3.5`), and `<virtualenv_name>` is a name of your choice for the virtual environment. For instance, it can be `py2.7` or `py3.7`.
|
||||
where `<version>` is the python version you want to use (select between `2.7` or `3.5`), and `<virtualenv_name>` is a name of your choice for the virtual environment. For instance, it can be `py2.7` or `py3.5`.
|
||||
|
||||
To deactivate the virtual environment, just type:
|
||||
```bash
|
||||
deactivate
|
||||
```
|
||||
|
||||
2. clone this repository and install the requirements and the setup.py
|
||||
2. clone this repository and install the requirements by executing the setup.py
|
||||
|
||||
In Python 2.7:
|
||||
```bash
|
||||
@@ -63,8 +63,8 @@ pip install -e . # this will install pyrobolearn as well as the required packag
|
||||
|
||||
Depending on your computer configuration and the python version you use, you might need to install also the following packages through `apt-get`:
|
||||
```bash
|
||||
sudo apt-get install python-tk # if python 2.7
|
||||
sudo apt-get install pytho3-tk # if python 3.5
|
||||
sudo apt install python-tk # if python 2.7
|
||||
sudo apt install python3-tk # if python 3.5
|
||||
```
|
||||
|
||||
## How to use it?
|
||||
|
||||
@@ -5,3 +5,10 @@ This folder contains storages / replay memories that are used in reinforcement l
|
||||
## What to look/check next?
|
||||
|
||||
Check the `algos` folder (especially the `algos/rl_algo` file) and the `samplers` folder.
|
||||
|
||||
## References
|
||||
|
||||
[1] "Reinforcement Learning for robots using neural networks", Lin, 1993
|
||||
[2] "Playing Atari with Deep Reinforcement Learning", Mnih et al., 2013
|
||||
[3] "Prioritized Experience Replay", Schaul, 2015
|
||||
[4] "Hindsight Experience Replay", Andrychowicz et al., 2017
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
|
||||
# import memory
|
||||
from storage import *
|
||||
from .storage import *
|
||||
|
||||
+123
-4
@@ -1,4 +1,123 @@
|
||||
# ER: Experience Replay (memory)
|
||||
# Refs:
|
||||
# 1. 'Reinforcement Learning for robots using neural networks', Lin, 1993
|
||||
# 2. 'Playing Atari with Deep Reinforcement Learning', Mnih et al., 2013
|
||||
#!/usr/bin/env python
|
||||
"""Provides the experience replay (ER) storage.
|
||||
|
||||
References:
|
||||
[1] "Reinforcement Learning for robots using neural networks", Lin, 1993
|
||||
[2] "Playing Atari with Deep Reinforcement Learning", Mnih et al., 2013
|
||||
"""
|
||||
|
||||
import random
|
||||
import torch
|
||||
|
||||
from pyrobolearn.storages.storage import Storage, ListStorage
|
||||
|
||||
|
||||
__author__ = "Brian Delhaisse"
|
||||
__copyright__ = "Copyright 2018, PyRoboLearn"
|
||||
__credits__ = ["Brian Delhaisse"]
|
||||
__license__ = "MIT"
|
||||
__version__ = "1.0.0"
|
||||
__maintainer__ = "Brian Delhaisse"
|
||||
__email__ = "briandelhaisse@gmail.com"
|
||||
__status__ = "Development"
|
||||
|
||||
|
||||
class ExperienceReplay(Storage):
|
||||
r"""Experience replay storage
|
||||
|
||||
The experience replay storage returns a transition tuple :math:`(s_t, a_t, s_{t+1}, r_t, d)`, where is :math:`s_t`
|
||||
is the state at time :math:`t`, :math:`a_t` is the action outputted by the policy in response to the state
|
||||
:math:`s_t`, :math:`s_{t+1}` is the next state returned by the environment due to the policy's action :math:`a_t`
|
||||
and the current state :math:`s_t`, :math:`r_t` is the reward signal returned by the environment, and :math:`d`
|
||||
is a boolean value that specifies if the task is over or not (i.e. if it has failed or succeeded).
|
||||
|
||||
The experience replay storage is often used in conjunction with off-policy RL algorithms.
|
||||
|
||||
The following code is inspired by [3] but modified such that it uses a PyTorch list storage.
|
||||
|
||||
References:
|
||||
[1] "Reinforcement Learning for robots using neural networks", Lin, 1993
|
||||
[2] "Playing Atari with Deep Reinforcement Learning", Mnih et al., 2013
|
||||
"""
|
||||
|
||||
def __init__(self, capacity=10000, device=None, dtype=torch.float):
|
||||
"""
|
||||
Initialize the Experience Replay Storage.
|
||||
|
||||
Args:
|
||||
capacity (int): maximum size of the experience replay storage.
|
||||
device (torch.device, str, None): the device to put the data on (e.g. `torch.device("cuda:0")` or
|
||||
`torch.device("cpu")`). If string, it can be 'cpu' or 'cuda'. If None, it will keep the original device
|
||||
to which the tensor is allocated.
|
||||
dtype (torch.dtype, None): convert the `torch.Tensor` to the specified data type. If None, it will keep
|
||||
the original dtype
|
||||
"""
|
||||
super(ExperienceReplay, self).__init__(device, dtype)
|
||||
self.capacity = capacity
|
||||
self.memory = ListStorage(device=device, dtype=dtype)
|
||||
self.position = 0
|
||||
|
||||
##############
|
||||
# Properties #
|
||||
##############
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
"""Return the memory's device instance."""
|
||||
return self.memory.device
|
||||
|
||||
@device.setter
|
||||
def device(self, device):
|
||||
"""Set the memory's device instance."""
|
||||
self.memory.device = device
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
"""Return the memory's data type."""
|
||||
return self.memory.dtype
|
||||
|
||||
@dtype.setter
|
||||
def dtype(self, dtype):
|
||||
"""Set the memory's data type."""
|
||||
self.memory.dtype = dtype
|
||||
|
||||
###########
|
||||
# Methods #
|
||||
###########
|
||||
|
||||
def push(self, *items):
|
||||
r"""
|
||||
Push new transition :math:`(s_t, a_t, s_{t+1}, r_t, d, \gamma)` in the experience replay.
|
||||
|
||||
Args:
|
||||
*items (list, tuple of torch.Tensor): transition tuple
|
||||
"""
|
||||
# add new item or update previous one
|
||||
if len(self.memory) < self.capacity:
|
||||
self.memory.append(items)
|
||||
else:
|
||||
self.memory[self.position] = items
|
||||
|
||||
# update head position (cyclic)
|
||||
self.position = (self.position + 1) % self.capacity
|
||||
|
||||
def to(self, device=None, dtype=None):
|
||||
"""
|
||||
Put all the tensors to the specified device and convert them to the specified data type.
|
||||
|
||||
Args:
|
||||
device (torch.device, str, None): the device to put the data on (e.g. `torch.device("cuda:0")` or
|
||||
`torch.device("cpu")`). If string, it can be 'cpu' or 'cuda'. If None, it will keep the original device
|
||||
to which the tensor is allocated.
|
||||
dtype (torch.dtype, None): convert the `torch.Tensor` to the specified data type. If None, it will keep
|
||||
the original dtype
|
||||
"""
|
||||
self.memory.to(device=device, dtype=dtype)
|
||||
|
||||
def sample(self, batch_size):
|
||||
"""Sample uniformly a batch from the experience replay."""
|
||||
return random.sample(self, batch_size)
|
||||
|
||||
|
||||
# alias
|
||||
ER = ExperienceReplay
|
||||
|
||||
@@ -1,2 +1,73 @@
|
||||
# HER: Hindsight Experience Replay
|
||||
# Ref: 'Hindsight Experience Replay', Andrychowicz et al., 2017
|
||||
#!/usr/bin/env python
|
||||
"""Provides the hindsight experience replay (HER) storage.
|
||||
|
||||
References:
|
||||
[1] "Hindsight Experience Replay", Andrychowicz et al., 2017
|
||||
"""
|
||||
|
||||
from pyrobolearn.storages.er import ExperienceReplay
|
||||
|
||||
__author__ = "Brian Delhaisse"
|
||||
__copyright__ = "Copyright 2018, PyRoboLearn"
|
||||
__credits__ = ["Brian Delhaisse"]
|
||||
__license__ = "MIT"
|
||||
__version__ = "1.0.0"
|
||||
__maintainer__ = "Brian Delhaisse"
|
||||
__email__ = "briandelhaisse@gmail.com"
|
||||
__status__ = "Development"
|
||||
|
||||
|
||||
class HindsightExperienceReplay(ExperienceReplay):
|
||||
r"""Hindsight Experience replay storage
|
||||
|
||||
One of the main challenges in RL is to shape the reward function such that the agent can successfully learned to
|
||||
perform the specified task. This often requires expert knowledge to engineer this reward function.
|
||||
To address this, the authors from [1] proposes to use a hindsight experience replay, which enables learning from
|
||||
sparse and binary rewards, and can be combined with any off-policy RL algorithms. This notably improves the
|
||||
sample efficiency.
|
||||
|
||||
In this setting, one or several goals have to be defined. They are concatenated with the state and feed to the
|
||||
policy and value approximators. Additionally, they are included in the transition tuple sampled from the
|
||||
experience replay storage.
|
||||
|
||||
|
||||
Pseudo-algo
|
||||
-----------
|
||||
|
||||
Pseudo-algorithm (taken from [1] and reproduce here for completeness)::
|
||||
1. Given:
|
||||
- an off-policy RL algorithm A (e.g. DQN, DDPG, NAF, SDQN)
|
||||
- a strategy S for sampling goals for replay (e.g. S(s_0, ..., s_T) = m(s_T))
|
||||
- a reward function r : S x A x G \rightarrow \mathbb{R} (e.g. r(s, a, g) = -[f_g(s) = 0])
|
||||
2. Initialize A (e.g. initialize neural networks)
|
||||
3. Initialize replay buffer R
|
||||
4. for episode = 1 to M do
|
||||
5. Sample a goal g and an initial state s0.
|
||||
6. for t = 0 to T - 1 do
|
||||
7. Sample an action at using the behavioral policy from A: a_t \leftarrow \pi_b([s_t,g])
|
||||
8. Execute the action a_t and observe a new state s_{t+1}
|
||||
9. end for
|
||||
10. for t = 0 to T - 1 do
|
||||
11. r_t := r(s_t, a_t, g)
|
||||
12. Store the transition ([s_t,g], a_t, r_t, [s_{t+1},g]) in R (standard experience replay)
|
||||
13. Sample a set of additional goals for replay G := S(current episode)
|
||||
14. for g' \in G do
|
||||
15. r' := r(s_t, a_t, g')
|
||||
16. Store the transition ([s_t,g'], a_t, r', [s_{t+1},g']) in R (HER)
|
||||
17. end for
|
||||
18. end for
|
||||
19. for t = 1 to N do
|
||||
20. Sample a minibatch B from the replay buffer R
|
||||
21. Perform one step of optimization using A and minibatch B
|
||||
22. end for
|
||||
23. end for
|
||||
|
||||
|
||||
References:
|
||||
[1] "Hindsight Experience Replay", Andrychowicz et al., 2017
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
# alias
|
||||
HER = HindsightExperienceReplay
|
||||
|
||||
@@ -1,2 +1,77 @@
|
||||
# PER: Prioritized Experience Replay
|
||||
# Ref: 'Prioritized Experience Replay', Schaul, 2015
|
||||
#!/usr/bin/env python
|
||||
"""Provides the prioritized experience replay (PER) storage.
|
||||
|
||||
The PER works by prioritizing transitions based on the magnitude of their TD error. In order to overcome over-fitting
|
||||
by sampling the same transitions, a stochastic sampling method is used based on a kind of softmax function on TD
|
||||
errors. In order to correct the bias induced by this sampling method, importance sampling weights (which are
|
||||
normalized for stability reasons) are used.
|
||||
|
||||
In summary, PER can be seen as a stochastic prioritization ER which uses importance sampling.
|
||||
|
||||
References:
|
||||
[1] "Prioritized Experience Replay", Schaul, 2015
|
||||
"""
|
||||
|
||||
from pyrobolearn.storages.storage import PriorityQueueStorage
|
||||
|
||||
__author__ = "Brian Delhaisse"
|
||||
__copyright__ = "Copyright 2018, PyRoboLearn"
|
||||
__credits__ = ["Brian Delhaisse"]
|
||||
__license__ = "MIT"
|
||||
__version__ = "1.0.0"
|
||||
__maintainer__ = "Brian Delhaisse"
|
||||
__email__ = "briandelhaisse@gmail.com"
|
||||
__status__ = "Development"
|
||||
|
||||
|
||||
class PrioritizedExperienceReplay(PriorityQueueStorage):
|
||||
r"""Prioritized Experience Replay storage
|
||||
|
||||
The PER works by prioritizing transitions based on the magnitude of their TD error. In order to overcome
|
||||
over-fitting by sampling the same transitions, a stochastic sampling method is used based on a kind of softmax
|
||||
function on TD errors. In order to correct the bias induced by this sampling method, importance sampling weights
|
||||
(which are normalized for stability reasons) are used.
|
||||
|
||||
In summary, PER can be seen as a stochastic prioritization ER which uses importance sampling.
|
||||
|
||||
There are 2 stochastic prioritization schemes used in [1].
|
||||
- proportional prioritization: :math:`p_i = |\delta_i| + \epsilon`, where :math:`\epsilon` is a small positive
|
||||
constant to avoid the transition to have a probability of 0.
|
||||
- rank-based prioritization: math:`p_i = \frac{1}{rank(i)}`, where rank(i) is the rank of transition i (that is
|
||||
they are i other keys in the priority queue that are smaller than the current key i) when the replay memory is
|
||||
sorted according to :math:`|\delta_i|`
|
||||
|
||||
|
||||
Pseudo-algo:
|
||||
-----------
|
||||
|
||||
Pseudo-algorithm (taken from [1] and reproduce here for completeness)::
|
||||
1. Input: minibatch k, step-size \eta, replay period K and size N, exponents a and b, budget T
|
||||
2. Initialize replay memory H = {}, \Delta = 0, p_1 = 1
|
||||
3. Observe s_0 and choose a_1 \sim \pi_\theta(s_0)
|
||||
4. for t = 0 to T-1 do
|
||||
5. choose action a_t \sim \pi_\theta(s_t)
|
||||
6. Observe s_{t+1}, r_t, \gamma_t
|
||||
7. Store transition (s_t, a_t, r_t, s_{t+1}, \gamma_t) in H with maximal priority p_t = max_{i<t} p_i
|
||||
8. if (t % K) = 0 then
|
||||
9. for j = 1 to k do
|
||||
10. Sample transition j \sim P(j) = \frac{ p_j^a }{ \sum_i p_i^a }
|
||||
11. Compute importance-sampling weight w_j = \frac{ (N P(j))^{-b} }{ \max_i w_i }
|
||||
12. Compute TD-error \delta_j = r_j + \gamma_j Q_{target}(s_j, argmax_a Q(s_j, a)) - Q(s_{j-1}, a_{j-1})
|
||||
13. Update transition priority p_j \leftarrow |\delta_j|
|
||||
14. Accumulate weight-change \Delta \leftarrow \Delta + w_j \delta_j \nabla_\theta Q(s_{j-1},a_{j-1})
|
||||
15. end for
|
||||
16. Update weights \theta \leftarrow \theta + \eta \Delta, reset \Delta = 0
|
||||
17. From time to time copy weights into target network \theta_{target} \leftarrow \theta
|
||||
18. end if
|
||||
19. end for
|
||||
|
||||
|
||||
References:
|
||||
[1] "Prioritized Experience Replay", Schaul, 2015
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
# alias
|
||||
PER = PrioritizedExperienceReplay
|
||||
|
||||
@@ -12,9 +12,9 @@ See Also:
|
||||
import collections
|
||||
import copy
|
||||
import pickle
|
||||
import queue
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler
|
||||
|
||||
from pyrobolearn import logger
|
||||
|
||||
@@ -118,6 +118,11 @@ class PyTorchStorage(Storage):
|
||||
if isinstance(item, torch.Tensor):
|
||||
logger.debug('setting tensor of size {} to {} with dtype={}'.format(item.size(), device, dtype))
|
||||
item = item.to(device=device, dtype=dtype)
|
||||
elif isinstance(item, np.ndarray):
|
||||
if item.dtype != object: # double, float, float16, int64, int32, and uint8
|
||||
item = torch.from_numpy(item).to(device=device, dtype=dtype)
|
||||
elif isinstance(item, (float, int, np.generic)):
|
||||
item = torch.tensor(item).to(device=device, dtype=dtype)
|
||||
elif isinstance(item, dict):
|
||||
for key, value in item.items():
|
||||
item[key] = self._to(value, device=device, dtype=dtype)
|
||||
@@ -127,8 +132,9 @@ class PyTorchStorage(Storage):
|
||||
value = self._to(value, device=device, dtype=dtype)
|
||||
item.add(value)
|
||||
elif isinstance(item, collections.Iterable):
|
||||
for idx, value in enumerate(item):
|
||||
item[idx] = self._to(value, device=device, dtype=dtype)
|
||||
item = [self._to(value, device=device, dtype=dtype) for value in item]
|
||||
# for idx, value in enumerate(item):
|
||||
# item[idx] = self._to(value, device=device, dtype=dtype)
|
||||
return item
|
||||
|
||||
@staticmethod
|
||||
@@ -232,6 +238,211 @@ class ListStorage(list, PyTorchStorage):
|
||||
iterable = self._to(iterable, device=self.device, dtype=self.dtype)
|
||||
super(ListStorage, self).append(iterable)
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
"""Set the specified value at the specified key."""
|
||||
super(ListStorage, self).__setitem__(key, self._to(value, device=self.device, dtype=self.dtype))
|
||||
|
||||
|
||||
class FIFOQueueStorage(queue.Queue, PyTorchStorage):
|
||||
r"""FIFO Queue Storage
|
||||
|
||||
FIFO queue storage (data structure) which allocates the given tensor(s) to the specified device.
|
||||
"""
|
||||
|
||||
def __init__(self, maxsize=0):
|
||||
"""Initialize the FIFO Queue storage.
|
||||
|
||||
Args:
|
||||
maxsize (int): maximum size of the queue. If :attr:`maxsize` is <= 0, the queue size is infinite.
|
||||
"""
|
||||
super(FIFOQueueStorage, self).__init__(maxsize)
|
||||
|
||||
def put(self, item, block=False, timeout=None):
|
||||
"""Put an item into the queue.
|
||||
|
||||
If optional args 'block' is true and 'timeout' is None (the default), block if necessary until a free slot
|
||||
is available. If 'timeout' is a non-negative number, it blocks at most 'timeout' seconds and raises
|
||||
the Full exception if no free slot was available within that time.
|
||||
Otherwise ('block' is false), put an item on the queue if a free slot is immediately available, else raise
|
||||
the Full exception ('timeout' is ignored in that case).
|
||||
"""
|
||||
if not self.full():
|
||||
item = self._to(item, device=self.device, dtype=self.dtype)
|
||||
super(FIFOQueueStorage, self).put(item, block=block, timeout=timeout)
|
||||
|
||||
def put_nowait(self, item):
|
||||
"""
|
||||
Put an item into the queue without blocking. Only enqueue the item if a free slot is immediately available.
|
||||
Otherwise raise the Full exception.
|
||||
"""
|
||||
item = self._to(item, device=self.device, dtype=self.dtype)
|
||||
super(FIFOQueueStorage, self).put_nowait(item)
|
||||
|
||||
def __len__(self):
|
||||
"""Return the size of the Queue."""
|
||||
return self.qsize()
|
||||
|
||||
def __iter__(self):
|
||||
"""Return the iterator object itself."""
|
||||
self.cnt = 0
|
||||
return self
|
||||
|
||||
def __next__(self): # only valid in Python 3
|
||||
"""Return the next item in the sequence."""
|
||||
if self.cnt < self.qsize():
|
||||
self.cnt += 1
|
||||
return self.queue[self.cnt-1]
|
||||
else:
|
||||
raise StopIteration
|
||||
|
||||
def next(self): # for Python 2
|
||||
"""Return the next item in the sequence."""
|
||||
return self.__next__()
|
||||
|
||||
|
||||
class LIFOQueueStorage(queue.LifoQueue, PyTorchStorage):
|
||||
r"""LIFO Queue Storage
|
||||
|
||||
LIFO queue storage (data structure) which allocates the given tensor(s) to the specified device.
|
||||
"""
|
||||
|
||||
def __init__(self, maxsize=0):
|
||||
"""Initialize the LIFO Queue storage.
|
||||
|
||||
Args:
|
||||
maxsize (int): maximum size of the queue. If :attr:`maxsize` is <= 0, the queue size is infinite.
|
||||
"""
|
||||
super(LIFOQueueStorage, self).__init__(maxsize)
|
||||
|
||||
def put(self, item, block=False, timeout=None):
|
||||
"""Put an item into the queue.
|
||||
|
||||
If optional args 'block' is true and 'timeout' is None (the default), block if necessary until a free slot
|
||||
is available. If 'timeout' is a non-negative number, it blocks at most 'timeout' seconds and raises
|
||||
the Full exception if no free slot was available within that time.
|
||||
Otherwise ('block' is false), put an item on the queue if a free slot is immediately available, else raise
|
||||
the Full exception ('timeout' is ignored in that case).
|
||||
"""
|
||||
if not self.full():
|
||||
item = self._to(item, device=self.device, dtype=self.dtype)
|
||||
super(LIFOQueueStorage, self).put(item, block=block, timeout=timeout)
|
||||
|
||||
def put_nowait(self, item):
|
||||
"""
|
||||
Put an item into the queue without blocking. Only enqueue the item if a free slot is immediately available.
|
||||
Otherwise raise the Full exception.
|
||||
"""
|
||||
item = self._to(item, device=self.device, dtype=self.dtype)
|
||||
super(LIFOQueueStorage, self).put_nowait(item)
|
||||
|
||||
def __len__(self):
|
||||
"""Return the size of the Queue."""
|
||||
return self.qsize()
|
||||
|
||||
def __iter__(self):
|
||||
"""Return the iterator object itself."""
|
||||
self.cnt = 0
|
||||
return self
|
||||
|
||||
def __next__(self): # only valid in Python 3
|
||||
"""Return the next item in the sequence."""
|
||||
if self.cnt < self.qsize():
|
||||
self.cnt += 1
|
||||
return self.queue[self.cnt-1]
|
||||
else:
|
||||
raise StopIteration
|
||||
|
||||
def next(self): # for Python 2
|
||||
"""Return the next item in the sequence."""
|
||||
return self.__next__()
|
||||
|
||||
|
||||
class PriorityQueueStorage(queue.PriorityQueue, PyTorchStorage):
|
||||
r"""Priority Queue Storage
|
||||
|
||||
Priority queue storage (data structure) which allocates the given tensor(s) to the specified device.
|
||||
|
||||
Note that `queue.PriorityQueue` is a thread-safe class that use the `heapq` module (which is initially not thread
|
||||
safe) under the hood.
|
||||
"""
|
||||
|
||||
def __init__(self, maxsize=0, ascending=True):
|
||||
"""Initialize the LIFO Queue storage.
|
||||
|
||||
Args:
|
||||
maxsize (int): maximum size of the queue. If :attr:`maxsize` is <= 0, the queue size is infinite.
|
||||
ascending (bool): if True, the item with the lowest priority will be the first one to be retrieved.
|
||||
"""
|
||||
super(PriorityQueueStorage, self).__init__(maxsize)
|
||||
self.ascending = ascending
|
||||
|
||||
def put(self, item, block=False, timeout=None):
|
||||
"""Put an item into the queue.
|
||||
|
||||
If optional args 'block' is true and 'timeout' is None (the default), block if necessary until a free slot
|
||||
is available. If 'timeout' is a non-negative number, it blocks at most 'timeout' seconds and raises
|
||||
the Full exception if no free slot was available within that time.
|
||||
Otherwise ('block' is false), put an item on the queue if a free slot is immediately available, else raise
|
||||
the Full exception ('timeout' is ignored in that case).
|
||||
"""
|
||||
if not self.full():
|
||||
if not isinstance(item, tuple) or len(item) != 2:
|
||||
raise TypeError("Expecting the item to be a tuple of length 2 with (priority number, data), instead "
|
||||
"got: {}".format(item))
|
||||
if self.ascending:
|
||||
item = (item[0], self._to(item[1], device=self.device, dtype=self.dtype))
|
||||
else:
|
||||
item = (-item[0], self._to(item[1], device=self.device, dtype=self.dtype))
|
||||
super(PriorityQueueStorage, self).put(item, block=block, timeout=timeout)
|
||||
|
||||
def put_nowait(self, item):
|
||||
"""
|
||||
Put an item into the queue without blocking. Only enqueue the item if a free slot is immediately available.
|
||||
Otherwise raise the Full exception.
|
||||
"""
|
||||
if not isinstance(item, tuple) or len(item) != 2:
|
||||
raise TypeError("Expecting the item to be a tuple of length 2 with (priority number, data), instead "
|
||||
"got: {}".format(item))
|
||||
if self.ascending:
|
||||
item = (item[0], self._to(item[1], device=self.device, dtype=self.dtype))
|
||||
else:
|
||||
item = (-item[0], self._to(item[1], device=self.device, dtype=self.dtype))
|
||||
super(PriorityQueueStorage, self).put_nowait(item)
|
||||
|
||||
def get(self, block=True, timeout=None):
|
||||
"""Remove and return an item from the queue.
|
||||
|
||||
If optional args 'block' is true and 'timeout' is None (the default), block if necessary until an item is
|
||||
available. If 'timeout' is a non-negative number, it blocks at most 'timeout' seconds and raises the Empty
|
||||
exception if no item was available within that time. Otherwise ('block' is false), return an item if one is
|
||||
immediately available, else raise the Empty exception ('timeout' is ignored in that case).
|
||||
"""
|
||||
item = super(PriorityQueueStorage, self).get(block=block, timeout=timeout)
|
||||
if not self.ascending:
|
||||
item = (-item[0], item[1])
|
||||
return item
|
||||
|
||||
def __len__(self):
|
||||
"""Return the size of the Queue."""
|
||||
return self.qsize()
|
||||
|
||||
def __iter__(self):
|
||||
"""Return the iterator object itself."""
|
||||
self.cnt = 0
|
||||
return self
|
||||
|
||||
def __next__(self): # only valid in Python 3
|
||||
"""Return the next item in the sequence."""
|
||||
if self.cnt < self.qsize():
|
||||
self.cnt += 1
|
||||
return self.queue[self.cnt-1]
|
||||
else:
|
||||
raise StopIteration
|
||||
|
||||
def next(self): # for Python 2
|
||||
"""Return the next item in the sequence."""
|
||||
return self.__next__()
|
||||
|
||||
|
||||
class SetStorage(set, PyTorchStorage):
|
||||
r"""PyTorch set storage
|
||||
|
||||
Reference in New Issue
Block a user