Unverified Commit bf8f05df authored by Muhammed Fatih BALIN's avatar Muhammed Fatih BALIN Committed by GitHub
Browse files

[GraphBolt][CUDA] Use deque instead of Queue for Bufferer. (#7050)

parent e117adac
"""Graph Bolt DataLoaders""" """Graph Bolt DataLoaders"""
from queue import Queue from collections import deque
import torch import torch
import torch.utils.data import torch.utils.data
...@@ -69,18 +69,18 @@ class Bufferer(dp.iter.IterDataPipe): ...@@ -69,18 +69,18 @@ class Bufferer(dp.iter.IterDataPipe):
raise ValueError( raise ValueError(
"'buffer_size' is required to be a positive integer." "'buffer_size' is required to be a positive integer."
) )
self.buffer = Queue(buffer_size) self.buffer = deque(maxlen=buffer_size)
def __iter__(self): def __iter__(self):
for data in self.datapipe: for data in self.datapipe:
if not self.buffer.full(): if len(self.buffer) < self.buffer.maxlen:
self.buffer.put(data) self.buffer.append(data)
else: else:
return_data = self.buffer.get() return_data = self.buffer.popleft()
self.buffer.put(data) self.buffer.append(data)
yield return_data yield return_data
while not self.buffer.empty(): while len(self.buffer) > 0:
yield self.buffer.get() yield self.buffer.popleft()
class Awaiter(dp.iter.IterDataPipe): class Awaiter(dp.iter.IterDataPipe):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment