def deseralize_pubkeys(self, data: bytearray, offset: Offset = None):
offset = offset or Offset()
num_keys = unpack_int(data, offset)
for _ in range(num_keys):
id = unpack_int(data, offset)
new_key = Key().deseralize_pubkey(data, offset)
self[id] = new_key
return self
def deseralize_keys(self, data: bytearray, offset: Offset = None):
offset = offset or Offset()
num_keys = unpack_int(data, offset)
for _ in range(num_keys):
new_key = Key()
new_key.deseralize_key(data, offset)
self.append(new_key)
return self
def decode(self, data):
offset = Offset()
self.src = unpack_int(data, offset)
self.dst = unpack_int(data, offset)
self.key = bytes(unpack_raw_data(data, offset))
self.data = unpack_raw_data(data, offset)
self.check()
return self
def decompress(self,
data: bytearray,
offset: Offset = None) -> torch.Tensor:
offset = offset or Offset()
torch.manual_seed(self.seed)
shape = unpack_tensor_shape(data, offset)
if len(shape) <= 1:
return unpack_tensor(data, offset)
else:
p = unpack_tensor(data, offset)
q = unpack_tensor(data, offset)
return torch.matmul(p, q.t()).view(shape)
def decompress(self,
data: bytearray,
offset: Offset = None) -> torch.Tensor:
offset = offset or Offset()
shape = unpack_tensor_shape(data, offset)
if len(shape) <= 1:
return unpack_tensor(data, offset)
else:
mean = unpack_float32(data, offset)
S = unpack_tensor(data, offset)
U = unpack_tensor(data, offset)
V = unpack_tensor(data, offset)
return (torch.mm(torch.mm(U, torch.diag(S)), V.t()) +
mean).view(shape)
def decompress(self, data: bytearray,
offset: Offset = None) -> torch.Tensor:
offset = offset or Offset()
# unpack tensor shape
shape = unpack_tensor_shape(data, offset)
# unpack sparse indices and value
tensor_indices = unpack_tensor(data, offset)
tensor_data = unpack_tensor(data, offset)
tensor = torch.zeros(int(np.prod(shape)))
tensor.index_copy_(0, tensor_indices, tensor_data)
tensor = tensor.reshape(list(shape))
return tensor
def decrypt(self, data: bytes, offset: Offset = None) -> bytes:
"""
throws ValueError exception when AES tag verification fails
"""
offset = offset or Offset()
nonce = unpack_raw_data(data, offset)
cipertext = unpack_raw_data(data, offset)
tag = unpack_raw_data(data, offset)
cipher = AES.new(self.key, AES.MODE_EAX, nonce=nonce)
plaintext = cipher.decrypt(cipertext)
cipher.verify(tag)
log.debug("The message is authentic")
return plaintext
def decompress(self, data: bytearray) -> [torch.Tensor]:
result = []
offset = Offset(0)
num_layers = unpack_int(data, offset)
if isinstance(self.compressors, _TensorCompressor):
compressors = [self.compressors] * num_layers
elif isinstance(self.compressors, list):
assert (len(self.compressors) == num_layers)
compressors = self.compressors
for i in range(num_layers):
result.append(compressors[i].decompress(data, offset))
return result
def decompress(self,
data: bytearray,
offset: Offset = None) -> torch.Tensor:
offset = offset or Offset()
# unpack original tensor shape
shape = unpack_tensor_shape(data, offset)
# unpack indices and data
indices = unpack_tensor(data, offset)
values = unpack_tensor(data, offset)
# re-construct tensor
tensor = torch.zeros(shape)
tensor.index_copy_(0, indices, values)
return tensor
def decompress(self,
data: bytearray,
offset: Offset = None) -> torch.Tensor:
offset = offset or Offset()
# unpack original tensor shape
shape = unpack_tensor_shape(data, offset)
# unpack compressed tensor
tensor_data = unpack_tensor(data, offset)
# restore random mask
k = math.ceil(self.ratio * np.prod(shape))
torch.manual_seed(self.seed)
mask = torch.randperm(int(np.prod(shape))).cuda() < k
idx = mask.int().sort(descending=True).indices
# restore tensor
tensor = torch.zeros(int(np.prod(shape)))
tensor.index_copy_(0, idx[:k].cpu(), tensor_data)
tensor = tensor.reshape(shape)
return tensor
def decompress(self, data: bytearray) -> [torch.Tensor]:
result = []
result_types = []
offset = Offset(0)
# get number of tensors
num_tensors = unpack_int(data, offset)
# re-construct the tensor list
for i in range(num_tensors):
dtype = bytes(unpack_fixed_width_data(data, offset, 1))
result_types.append(self.type_reverse_map[dtype])
shape = unpack_tensor_shape(data, offset)
result.append(torch.zeros(shape))
# restore compressed data
compressed_data = self.compressors.decompress(data, offset)
torch.nn.utils.vector_to_parameters(compressed_data, result)
# restore types
for i in range(len(result)):
result[i] = result[i].to(dtype=result_types[i])
return result
def process_step_weights(self, socket, data):
client_id = self.clients_socket_id_map[socket]
self.clients_step_weights[client_id] = unpack_float32(data, Offset())
self.sem_clients_step_weights.release()
def decompress(self,
data: bytearray,
offset: Offset = None) -> torch.Tensor:
offset = offset or Offset()
return unpack_tensor(data, offset)
def deseralize_key(self, data, offset:Offset = None):
offset = offset or Offset()
self.key = bytes(unpack_raw_data(data, offset))
return self
def decompress(self,
data: bytearray,
offset: Offset = None) -> torch.Tensor:
offset = offset or Offset()
raise NotImplementedError("Usage of abstract method.")
def decompress(self,
data: bytearray,
offset: Offset = None) -> torch.Tensor:
offset = offset or Offset()
raise NotImplementedError
def deseralize_pubkey(self, data, offset: Offset = None):
offset = offset or Offset()
self.pub_key = RSA.import_key(bytes(unpack_raw_data(data, offset)))
self.cipher = PKCS1_OAEP.new(self.pub_key, hashAlgo=SHA256)
# we cannot decrypt without a priv key, so not setting self.decrypt
return self