espnet2/utils/sized_dict.py

import collections
import sys

from torch import multiprocessing


def get_size(obj, seen=None):
    """Recursively finds size of objects

    Taken from https://github.com/bosswissam/pysize

    """

    size = sys.getsizeof(obj)
    if seen is None:
        seen = set()

    obj_id = id(obj)
    if obj_id in seen:
        return 0

    # Important mark as seen *before* entering recursion to gracefully handle
    # self-referential objects
    seen.add(obj_id)

    if isinstance(obj, dict):
        size += sum([get_size(v, seen) for v in obj.values()])
        size += sum([get_size(k, seen) for k in obj.keys()])
    elif hasattr(obj, "__dict__"):
        size += get_size(obj.__dict__, seen)
    elif isinstance(obj, (list, set, tuple)):
        size += sum([get_size(i, seen) for i in obj])

    return size


class SizedDict(collections.abc.MutableMapping):
    def __init__(self, shared: bool = False, data: dict = None):
        if data is None:
            data = {}

        if shared:
            # NOTE(kamo): Don't set manager as a field because Manager, which includes
            # weakref object, causes following error with method="spawn",
            # "TypeError: can't pickle weakref objects"
            self.cache = multiprocessing.Manager().dict(**data)
        else:
            self.manager = None
            self.cache = dict(**data)
        self.size = 0

    def __setitem__(self, key, value):
        if key in self.cache:
            self.size -= get_size(self.cache[key])
        else:
            self.size += sys.getsizeof(key)
        self.size += get_size(value)
        self.cache[key] = value

    def __getitem__(self, key):
        return self.cache[key]

    def __delitem__(self, key):
        self.size -= get_size(self.cache[key])
        self.size -= sys.getsizeof(key)
        del self.cache[key]

    def __iter__(self):
        return iter(self.cache)

    def __contains__(self, key):
        return key in self.cache

    def __len__(self):
        return len(self.cache)