45 lines
1.1 KiB
Python
45 lines
1.1 KiB
Python
import os
|
|
import sys
|
|
import numpy as np
|
|
from torch.utils.data import DataLoader
|
|
|
|
|
|
|
|
class Accumulator():
|
|
def __init__(self):
|
|
self.dict = dict()
|
|
|
|
def update(self, dict):
|
|
for key in dict.keys():
|
|
if key in self.dict.keys():
|
|
self.dict[key].append(dict[key])
|
|
else:
|
|
self.dict[key] = [dict[key]]
|
|
|
|
def get(self):
|
|
result = dict()
|
|
for key in self.dict.keys():
|
|
result[key] = np.mean(self.dict[key])
|
|
return result
|
|
|
|
def reset(self):
|
|
self.dict = dict()
|
|
|
|
|
|
class InfiniteDataLoader(DataLoader):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
# Initialize an iterator over the dataset.
|
|
self.dataset_iterator = super().__iter__()
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
try:
|
|
batch = next(self.dataset_iterator)
|
|
except StopIteration:
|
|
# Dataset exhausted, use a new fresh iterator.
|
|
self.dataset_iterator = super().__iter__()
|
|
batch = next(self.dataset_iterator)
|
|
return batch |