def push(self, model: nn.Module, pull_on_fail=True):
"""
Try to push a model to the ordered server, if failed, the newest
model will be automatically pulled and its parameters will be
assigned to ``model``. Gradients will not be cleared.
Args:
model: Model to push.
pull_on_fail: Pull the newest parameters if push failed.
"""
if not hasattr(model, "pp_version"):
model.pp_version = 0
copied_model_params = deepcopy(model.state_dict())
for k, v in copied_model_params.items():
copied_model_params[k] = v.to("cpu")
if not self.o_server.push(
self.model_name, copied_model_params,
version=model.pp_version + 1, prev_version=model.pp_version
):
if pull_on_fail:
result = self.o_server.pull(self.model_name)
if result is None: # pragma: no cover
raise RuntimeError("Pull failed, this should not happen.")
st_dict, version = result
prep_load_state_dict(model, st_dict)
model.pp_version = version
return False
else:
model.pp_version += 1
return True
def test_prep_load_state_dict(pytestconfig):
model = t.nn.Linear(100, 100)
model2 = t.nn.Linear(100, 100).to(pytestconfig.getoption("gpu_device"))
state_dict = model2.state_dict()
prep_load_state_dict(model, state_dict)
assert t.all(model.weight == model2.weight.to("cpu"))
assert t.all(model.bias == model2.bias.to("cpu"))
def pull(self, model: nn.Module):
"""
Pull the newest model.
Args:
model: Model to push.
"""
params = self.server.pull("_push_pull_grad_managed_model")
prep_load_state_dict(model, params)
def pull(self, model: nn.Module):
"""
Pull the newest model. Its gradients will be cleared.
Args:
model: Model to push.
"""
model.zero_grad()
params = self.o_server.pull(self.model_name)
if params is not None:
# params could be None if the master reducer has't performed
# a single reduction operation yet
prep_load_state_dict(model, params[0])
def pull(self, model: nn.Module):
"""
Pull the newest state dict of your model and update its parameters
and ``pp_version``.
Args:
model: Model to pull.
"""
result = self.server.pull(self.MODEL_NAME)
if result is None:
return
st_dict, version = result
if not hasattr(model, "pp_version") or model.pp_version < version:
prep_load_state_dict(model, st_dict)
model.pp_version = version
def pull(self, model: nn.Module):
"""
Pull the newest state dict of your model and update its parameters
and ``pp_version``. Gradients will not be cleared.
Args:
model: Model to pull.
"""
result = self.o_server.pull(self.model_name)
if result is None: # pragma: no cover
return False
st_dict, version = result
if not hasattr(model, "pp_version") or model.pp_version < version:
prep_load_state_dict(model, st_dict)
model.pp_version = version
return True