simple_demo.py 1.2 KB
Newer Older
Z
Zeyu Chen 已提交
1 2 3
import paddle.fluid as fluid
import paddlehub as hub

Z
Zeyu Chen 已提交
4
# Step1: Select pre-trained model
Z
Zeyu Chen 已提交
5 6 7
module = hub.Module(name="ernie")
inputs, outputs, program = module.context(trainable=True, max_seq_len=128)

Z
Zeyu Chen 已提交
8
# Step2: Prepare Dataset and DataReader
Z
Zeyu Chen 已提交
9 10 11 12
dataset = hub.dataset.ChnSentiCorp()
reader = hub.reader.ClassifyReader(
    dataset=dataset, vocab_path=module.get_vocab_path(), max_seq_len=128)

Z
Zeyu Chen 已提交
13
# Step3: Construct transfer learning task
Z
Zeyu Chen 已提交
14 15 16
with fluid.program_guard(program):
    label = fluid.layers.data(name="label", shape=[1], dtype='int64')
    pooled_output = outputs["pooled_output"]
Z
Zeyu Chen 已提交
17
    cls_task = hub.create_text_cls_task(
Z
Zeyu Chen 已提交
18 19
        feature=pooled_output, label=label, num_classes=dataset.num_labels)

Z
Zeyu Chen 已提交
20
# Step4: Setup config then start finetune
Z
Zeyu Chen 已提交
21 22 23
strategy = hub.AdamWeightDecayStrategy(learning_rate=5e-5, weight_decay=0.01)

config = hub.RunConfig(
Z
Zeyu Chen 已提交
24 25 26 27 28
    use_cuda=True,
    checkpoint_dir="./ckpt",
    num_epoch=3,
    batch_size=32,
    strategy=strategy)
Z
Zeyu Chen 已提交
29 30 31 32 33 34 35 36

feed_list = [
    inputs["input_ids"].name, inputs["position_ids"].name,
    inputs["segment_ids"].name, inputs["input_mask"].name, label.name
]

hub.finetune_and_eval(
    task=cls_task, data_reader=reader, feed_list=feed_list, config=config)