elif pretrained_bert_parameters:
# only load BertModel parameters
nlp.utils.load_parameters(bert,
pretrained_bert_parameters,
ctx=ctx,
ignore_extra=True,
cast_dtype=True)
net.span_classifier.initialize(init=mx.init.Normal(0.02), ctx=ctx)
elif pretrained:
# only load BertModel parameters
net.span_classifier.initialize(init=mx.init.Normal(0.02), ctx=ctx)
else:
# no checkpoint is loaded
net.initialize(init=mx.init.Normal(0.02), ctx=ctx)
net.hybridize(static_alloc=True)
loss_function = BertForQALoss()
loss_function.hybridize(static_alloc=True)
def train():
"""Training function."""
segment = 'train' if not args.debug else 'dev'
log.info('Loading %s data...', segment)
if version_2:
train_data = SQuAD(segment, version='2.0')
else:
train_data = SQuAD(segment, version='1.1')
if args.debug:
sampled_data = [train_data[i] for i in range(1000)]
elif pretrained_bert_parameters:
# only load BertModel parameters
nlp.utils.load_parameters(bert,
pretrained_bert_parameters,
ctx=ctx,
ignore_extra=True,
cast_dtype=True)
net.span_classifier.initialize(init=mx.init.Normal(0.02), ctx=ctx)
elif pretrained:
# only load BertModel parameters
net.span_classifier.initialize(init=mx.init.Normal(0.02), ctx=ctx)
else:
# no checkpoint is loaded
net.initialize(init=mx.init.Normal(0.02), ctx=ctx)
net.hybridize(static_alloc=True)
loss_function = BertForQALoss()
loss_function.hybridize(static_alloc=True)
if deploy:
logging.info(
'load symbol file directly as SymbolBlock for model deployment')
net = mx.gluon.SymbolBlock.imports(
'{}-symbol.json'.format(args.model_prefix),
['data0', 'data1', 'data2'],
'{}-0000.params'.format(args.model_prefix))
net.hybridize(static_alloc=True, static_shape=True)
# calibration config
only_calibration = args.only_calibration
