def testCompressUpdate(self):
params, info = self._construct_info()
params0 = copy.deepcopy(params)
apply_gradients([p.grad for p in params], params0)
optimizer = RMSprop(params0, lr=0.1, momentum=0.5)
optimizer.step()
params1 = copy.deepcopy(params)
apply_gradients([p.grad for p in params], params1)
optimizer1 = RMSprop(params1, lr=0.1, momentum=0.5)
optimizer1.step()
assertAllClose(params0[1], params1[1])
assertAllClose(params0[2], params1[2])
assertAllClose(params0[0], params1[0])
info['var_old'] = params1[0]
optimizer1.compress_mask(info, verbose=True)
optimizer1.compress_drop({'var_old': params1[2], 'type': 'variable'})
info['mask_hook'](info['var_new'], info['var_old'], info['mask'])
params1[0] = info['var_new']
params1[0].grad = params0[0].grad.data[info['mask']]
optimizer1.step() # params1[2] not updated
assertAllClose(params0[2], params1[2])
optimizer.step()
assertAllClose(params0[1], params1[1])
assertAllClose(params0[0][info['mask']], params1[0])
def testCompressWithoutState(self):
parameters, info = self._construct_info()
optimizer = RMSprop(parameters)
optimizer.compress_mask(info, verbose=False)
self.assertFalse(check_tensor_in(info['var_new'], optimizer.state))
self.assertTrue(
check_tensor_in(info['var_new'],
optimizer.param_groups[0]['params']))
self.assertFalse(check_tensor_in(parameters[0], optimizer.state))
self.assertFalse(
check_tensor_in(parameters[0],
optimizer.param_groups[0]['params']))