def test_doubleMODELcompression_WITH_memory(self):
     params = SGDDoubleModelCompressionWithMem().define(
         n_dimensions=DIM, nb_devices=1, quantization_param=10)
     params.learning_rate = 0.5
     workers = [Worker(0, params)]
     workers[0].set_data(x, y)
     workers[0].cost_model.L = workers[0].cost_model.local_L
     update = ArtemisUpdate(params, workers)
     artificial_l = ones_tensor.clone().detach()
     update.l = artificial_l.clone().detach()
     new_w = update.compute(w, 2, 2)
     # Check that gradients have been updated.
     self.assertFalse(torch.equal(update.g, zero_tensor))
     self.assertFalse(torch.equal(update.v, zero_tensor))
     self.assertFalse(torch.equal(update.h, zero_tensor))
     # Check that l has been updated.
     self.assertFalse(torch.equal(update.l, artificial_l))
     # Check that correct value has been compressed
     self.assertTrue(
         torch.equal(update.value_to_compress, new_w - artificial_l))
 def test_doubleGRADIENTcompression_WITH_additional_memory(self):
     params = DoreVariant().define(n_dimensions=DIM,
                                   nb_devices=1,
                                   quantization_param=10)
     params.learning_rate = 0.5
     workers = [Worker(0, params)]
     workers[0].set_data(x, y)
     workers[0].cost_model.L = workers[0].cost_model.local_L
     update = ArtemisUpdate(params, workers)
     artificial_l = ones_tensor.clone().detach()
     # We artificially set different memory to check that it has impact on update computation.
     update.l = artificial_l.clone().detach()
     update.compute(w, 2, 2)
     # Check that gradients have been updated.
     self.assertFalse(torch.equal(update.g, zero_tensor))
     self.assertFalse(torch.equal(update.v, zero_tensor))
     self.assertFalse(torch.equal(update.h, zero_tensor))
     # Check that l has been updated.
     self.assertFalse(torch.equal(update.l, artificial_l))
     # Check that correct value has been compressed
     self.assertTrue(
         torch.equal(update.value_to_compress, update.g - artificial_l))