def test_load_from_chkpt(config_path, chkpt_path, chkpt_task, custom_ops):
"""
Compare the model loaded into our popart model against the modified
PyTorch model:
- Load tf weights into BERT using torch impl -> run fwd model
- Load tf weights into BERT using popart impl -> run fwd model
- Compare output tensors
"""
# Vocab-size override is not required, but allows the test to run more quickly
config = load_bert_config_tf(config_path,
override_vocab=9728,
chkpt_task=chkpt_task)
builder = popart.Builder(opsets={
"ai.onnx": 9,
"ai.onnx.ml": 1,
"ai.graphcore": 1
})
# Load Torch version
TorchModel = TorchModelPreTraining if chkpt_task == "PRETRAINING" else TorchModelSquad
torch_model = TorchModel(
TorchBertConfig(
config.vocab_length,
config.hidden_size,
num_hidden_layers=config.num_layers,
num_attention_heads=config.attention_heads,
intermediate_size=config.ff_size,
hidden_act="relu",
max_position_embeddings=config.max_positional_length,
layer_norm_eps=config.layer_norm_eps,
mask_tokens=config.mask_tokens,
))
torch_model = load_tf_weights_in_bert(torch_model, config, chkpt_path)
# Load Popart model
sequence_info = popart.TensorInfo(
"UINT32", [config.batch_size * config.sequence_length])
indices = builder.addInputTensor(sequence_info)
positions = builder.addInputTensor(sequence_info)
segments = builder.addInputTensor(sequence_info)
popart_model, proto, output = load_model_from_tf(chkpt_path,
True,
config,
indices,
positions,
segments,
chkpt_task,
builder=builder)
run_models(config, proto, indices, positions, segments, output,
popart_model, torch_model)
def test_load_from_frozen(config_path, chkpt_path, frozen_path, custom_ops):
# Vocab-size override is not required, but allows the test to run more quickly
config = load_bert_config_tf(config_path, override_vocab=9728)
builder = popart.Builder(
opsets={"ai.onnx": 9, "ai.onnx.ml": 1, "ai.graphcore": 1}
)
# Load Torch version
torch_model = TorchModel(
TorchBertConfig(
config.vocab_length,
config.hidden_size,
num_hidden_layers=config.num_layers,
num_attention_heads=config.attention_heads,
intermediate_size=config.ff_size,
hidden_act="relu",
max_position_embeddings=config.max_positional_length,
layer_norm_eps=config.layer_norm_eps,
mask_tokens=config.mask_tokens,
)
)
torch_model = load_tf_weights_in_bert(torch_model, config, chkpt_path)
# Load Popart model
sequence_info = popart.TensorInfo(
"UINT32", [config.batch_size * config.sequence_length])
indices = builder.addInputTensor(sequence_info)
positions = builder.addInputTensor(sequence_info)
segments = builder.addInputTensor(sequence_info)
popart_model, proto, output = load_model_from_tf(
frozen_path, False, config, indices, positions, segments, builder=builder
)
run_models(config, proto, indices, positions, segments, output, popart_model, torch_model)
def test_load_from_chkpt(config_path, chkpt_path, custom_ops):
"""
Compare the model loaded into our popart model against the modified
PyTorch model:
- Load tf weights into BERT using torch impl -> run fwd model
- Load tf weights into BERT using popart impl -> run fwd model
- Compare output tensors
"""
config = load_bert_config_tf(config_path)
builder = popart.Builder(opsets={
"ai.onnx": 9,
"ai.onnx.ml": 1,
"ai.graphcore": 1
})
# Load Torch version
torch_model = TorchModel(
TorchBertConfig(
config.vocab_length,
config.hidden_size,
num_hidden_layers=config.num_layers,
num_attention_heads=config.attention_heads,
intermediate_size=config.ff_size,
hidden_act="relu",
max_position_embeddings=config.max_positional_length,
layer_norm_eps=config.layer_norm_eps,
mask_tokens=config.mask_tokens,
))
torch_model.eval()
torch_model = load_tf_weights_in_bert(torch_model, config, chkpt_path)
# Load Popart model
sequence_info = popart.TensorInfo(
"INT32", [config.batch_size * config.sequence_length])
indices = builder.addInputTensor(sequence_info)
positions = builder.addInputTensor(sequence_info)
popart_model, proto, output = load_from_tf(chkpt_path,
True,
config,
indices,
positions,
builder=builder)
# Run the models
popart_inputs = {
indices:
np.random.randint(0, config.vocab_length,
(config.batch_size * config.sequence_length)).astype(
np.int32),
positions:
np.random.randint(
0,
config.sequence_length,
(config.batch_size * config.sequence_length),
).astype(np.int32),
}
torch_inputs = {
"input_ids":
popart_inputs[indices].reshape(config.batch_size,
config.sequence_length),
"position_ids":
popart_inputs[positions].reshape(config.batch_size,
config.sequence_length),
}
torch_outputs = run_fwd_model(torch_inputs, torch_model)
popart_outputs, post_proto = run_py(
proto,
popart_inputs,
output,
ipus=math.ceil(config.num_layers / config.layers_per_ipu) + 1,
)
check_tensors(torch_outputs, popart_outputs)
print("Test succeeded")