def testToyBertCheckpointFrozenWeights():
# Common setup
seed = 1
total_steps = 10
torch.manual_seed(seed)
onnxruntime.set_seed(seed)
opts = orttrainer.ORTTrainerOptions({
'debug': {
'deterministic_compute': True
},
'utils': {
'frozen_weights':
['bert.encoder.layer.0.attention.self.value.weight']
}
})
# Create ORTTrainer and save initial state in a dict
model = load_bert_onnx_model()
model_desc = bert_model_description()
optim_config = optim.LambConfig()
trainer = orttrainer.ORTTrainer(model,
model_desc,
optim_config,
options=opts)
# Train for a few steps
for i in range(total_steps):
sample_input = generate_random_input_from_model_desc(model_desc, seed)
_ = trainer.train_step(*sample_input)
sample_input = generate_random_input_from_model_desc(
model_desc, seed + total_steps + 1)
# Evaluate once to get a base loss
loss = trainer.eval_step(*sample_input)
# Save checkpoint
state_dict = checkpoint.experimental_state_dict(trainer)
# Load previous state into another instance of ORTTrainer
model2 = load_bert_onnx_model()
model_desc2 = bert_model_description()
optim_config2 = optim.LambConfig()
trainer2 = orttrainer.ORTTrainer(model2,
model_desc2,
optim_config2,
options=opts)
checkpoint.experimental_load_state_dict(trainer2, state_dict)
# Evaluate once to get a base loss
ckpt_loss = trainer2.eval_step(*sample_input)
# Must match as both trainers have the same dict state
assert_allclose(loss.cpu(), ckpt_loss.cpu())
loaded_state_dict = checkpoint.experimental_state_dict(trainer2)
assert state_dict.keys() == loaded_state_dict.keys()
def testToyBertCheckpointBasic():
# Common setup
seed = 1
torch.manual_seed(seed)
onnxruntime.set_seed(seed)
optim_config = optim.LambConfig()
opts = orttrainer.ORTTrainerOptions(
{'debug': {
'deterministic_compute': True
}})
# Create ORTTrainer and save initial state in a dict
model = load_bert_onnx_model()
model_desc = bert_model_description()
trainer = orttrainer.ORTTrainer(model,
model_desc,
optim_config,
options=opts)
sd = checkpoint.experimental_state_dict(trainer)
## All initializers must be present in the state_dict
## when the specified model for ORTTRainer is an ONNX model
for param in trainer._onnx_model.graph.initializer:
assert param.name in sd
## Modify one of the state values and load into ORTTrainer
sd['bert.encoder.layer.0.attention.output.LayerNorm.weight'] += 10
checkpoint.experimental_load_state_dict(trainer, sd)
## Save a checkpoint
ckpt_dir = _test_helpers._get_name("ort_ckpt")
checkpoint.experimental_save_checkpoint(trainer, ckpt_dir,
'bert_toy_save_test')
del trainer
del model
# Create a new ORTTrainer and load the checkpoint from previous ORTTrainer
model2 = load_bert_onnx_model()
model_desc2 = bert_model_description()
trainer2 = orttrainer.ORTTrainer(model2,
model_desc2,
optim_config,
options=opts)
checkpoint.experimental_load_checkpoint(trainer2, ckpt_dir,
'bert_toy_save_test')
loaded_sd = checkpoint.experimental_state_dict(trainer2)
# Assert whether original state and the one loaded from checkpoint matches
for k, v in loaded_sd.items():
assert torch.all(torch.eq(v, sd[k]))