mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2024-12-25 07:21:13 -08:00
112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
import multiprocessing
|
|
import pickle
|
|
import struct
|
|
from core.joblib import Subprocessor
|
|
|
|
class MPSharedList():
|
|
"""
|
|
Provides read-only pickled list of constant objects via shared memory aka 'multiprocessing.Array'
|
|
Thus no 4GB limit for subprocesses.
|
|
|
|
supports list concat via + or sum()
|
|
"""
|
|
|
|
def __init__(self, obj_list):
|
|
if obj_list is None:
|
|
self.obj_counts = None
|
|
self.table_offsets = None
|
|
self.data_offsets = None
|
|
self.sh_bs = None
|
|
else:
|
|
obj_count, table_offset, data_offset, sh_b = MPSharedList.bake_data(obj_list)
|
|
|
|
self.obj_counts = [obj_count]
|
|
self.table_offsets = [table_offset]
|
|
self.data_offsets = [data_offset]
|
|
self.sh_bs = [sh_b]
|
|
|
|
def __add__(self, o):
|
|
if isinstance(o, MPSharedList):
|
|
m = MPSharedList(None)
|
|
m.obj_counts = self.obj_counts + o.obj_counts
|
|
m.table_offsets = self.table_offsets + o.table_offsets
|
|
m.data_offsets = self.data_offsets + o.data_offsets
|
|
m.sh_bs = self.sh_bs + o.sh_bs
|
|
return m
|
|
elif isinstance(o, int):
|
|
return self
|
|
else:
|
|
raise ValueError(f"MPSharedList object of class {o.__class__} is not supported for __add__ operator.")
|
|
|
|
def __radd__(self, o):
|
|
return self+o
|
|
|
|
def __len__(self):
|
|
return sum(self.obj_counts)
|
|
|
|
def __getitem__(self, key):
|
|
obj_count = sum(self.obj_counts)
|
|
if key < 0:
|
|
key = obj_count+key
|
|
if key < 0 or key >= obj_count:
|
|
raise ValueError("out of range")
|
|
|
|
for i in range(len(self.obj_counts)):
|
|
|
|
if key < self.obj_counts[i]:
|
|
table_offset = self.table_offsets[i]
|
|
data_offset = self.data_offsets[i]
|
|
sh_b = self.sh_bs[i]
|
|
break
|
|
key -= self.obj_counts[i]
|
|
|
|
sh_b = memoryview(sh_b).cast('B')
|
|
|
|
offset_start, offset_end = struct.unpack('<QQ', sh_b[ table_offset + key*8 : table_offset + (key+2)*8].tobytes() )
|
|
|
|
return pickle.loads( sh_b[ data_offset + offset_start : data_offset + offset_end ].tobytes() )
|
|
|
|
def __iter__(self):
|
|
for i in range(self.__len__()):
|
|
yield self.__getitem__(i)
|
|
|
|
@staticmethod
|
|
def bake_data(obj_list):
|
|
if not isinstance(obj_list, list):
|
|
raise ValueError("MPSharedList: obj_list should be list type.")
|
|
|
|
obj_count = len(obj_list)
|
|
|
|
if obj_count != 0:
|
|
obj_pickled_ar = [pickle.dumps(o, 4) for o in obj_list]
|
|
|
|
table_offset = 0
|
|
table_size = (obj_count+1)*8
|
|
data_offset = table_offset + table_size
|
|
data_size = sum([len(x) for x in obj_pickled_ar])
|
|
|
|
sh_b = multiprocessing.RawArray('B', table_size + data_size)
|
|
#sh_b[0:8] = struct.pack('<Q', obj_count)
|
|
sh_b_view = memoryview(sh_b).cast('B')
|
|
|
|
offset = 0
|
|
|
|
sh_b_table = bytes()
|
|
offsets = []
|
|
|
|
offset = 0
|
|
for i in range(obj_count):
|
|
offsets.append(offset)
|
|
offset += len(obj_pickled_ar[i])
|
|
offsets.append(offset)
|
|
|
|
sh_b_view[table_offset:table_offset+table_size] = struct.pack( '<'+'Q'*len(offsets), *offsets )
|
|
|
|
for i, obj_pickled in enumerate(obj_pickled_ar):
|
|
offset = data_offset+offsets[i]
|
|
sh_b_view[offset:offset+len(obj_pickled)] = obj_pickled_ar[i]
|
|
|
|
return obj_count, table_offset, data_offset, sh_b
|
|
return 0, 0, 0, None
|
|
|