Add support for partial checkpoint restore

[terrykong]

This change allows partial checkpoint restores to be configurable, which is often needed when fine-tuning from a pretrained checkpoint. You can add the following code to your gin config to restore only a subset of the model variables from the given checkpoint:

utils.RestoreCheckpointConfig:
  fallback_to_scratch = True
  strict = False
  state_transformation_fns = [@state_transform]

where state_transform determines which parameters to restore. For example, to restore only model weights and reset optimizer states, we can use the following transformation function:

def just_states_transform(checkpoint, opt_state):
    return state_utils.apply_assignment_map(checkpoint, opt_state,
        assignment_map=[(r'state.*', None)])

Mapping all variables beginning with state. (i.e. the optimizer states) to None causes them to be ignored during checkpoint restore. All variables beginning with target. (i.e. the model weights) will be restored.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.