class TransferableWeights(metaclass=segment_transfer_enabled()):
def __init__(self, weights, cls, *args, **kwargs):
self._weights = weights
self._cls = cls
if args:
self._args = args
if kwargs:
self._kwargs = kwargs
def with_degree(self, degree):
setattr(self, "_degree", degree)
return self
def get_degree(self, default=None):
return getattr(self, "_degree", default)
@property
def unboxed(self):
return self._weights
@property
def weights(self):
if not hasattr(self, "_args") and not hasattr(self, "_kwargs"):
return self._cls(self._weights)
else:
args = self._args if hasattr(self, "_args") else ()
kwargs = self._kwargs if hasattr(self, "_kwargs") else {}
return self._cls(self._weights, *args, **kwargs)
class Weights(metaclass=segment_transfer_enabled()):
def __init__(self, l):
self._weights = l
def for_remote(self):
return TransferableWeights(self._weights, self.__class__)
@property
def unboxed(self):
return self._weights
@abc.abstractmethod
def map_values(self, func, inplace):
pass
@abc.abstractmethod
def binary_op(self, other, func, inplace):
pass
@abc.abstractmethod
def axpy(self, a, y):
pass
def decrypted(self, cipher: Encrypt, inplace=True):
return self.map_values(cipher.decrypt, inplace=inplace)
def encrypted(self, cipher: Encrypt, inplace=True):
return self.map_values(cipher.encrypt, inplace=inplace)
def __imul__(self, other):
return self.map_values(lambda x: x * other, inplace=True)
def __mul__(self, other):
return self.map_values(lambda x: x * other, inplace=False)
def __iadd__(self, other):
return self.binary_op(other, operator.add, inplace=True)
def __add__(self, other):
LOGGER.debug("In binary_op0, _w: {}".format(self._weights))
return self.binary_op(other, operator.add, inplace=False)
def __isub__(self, other):
return self.binary_op(other, operator.sub, inplace=True)
def __sub__(self, other):
return self.binary_op(other, operator.sub, inplace=False)
def __truediv__(self, other):
return self.map_values(lambda x: x / other, inplace=False)
def __itruediv__(self, other):
return self.map_values(lambda x: x / other, inplace=True)
class TransferableGradients(metaclass=segment_transfer_enabled(max_part_size=FRAGMENT_24M)):
def __init__(self, gradients, cls, *args, **kwargs):
self._gradients = gradients
self._cls = cls
if args:
self._args = args
if kwargs:
self._kwargs = kwargs
@property
def unboxed(self):
return self._gradients
@property
def gradients(self):
if not hasattr(self, "_args") and not hasattr(self, "_kwargs"):
return self._cls(self._gradients)
else:
args = self._args if hasattr(self, "_args") else ()
kwargs = self._kwargs if hasattr(self, "_kwargs") else {}
return self._cls(self._gradients, *args, **kwargs)
class BigObjectTransfer(metaclass=segment_transfer_enabled()):
def __init__(self, data):
self._obj = data
def get_data(self):
return self._obj