Skip to main content

BERT for Multi-task Learning

Project description

Bert for Multi-task Learning

python tensorflow PyPI version PyPI license


Note: Since 0.4.0, tf version >= 2.1 is required.


pip install bert-multitask-learning

What is it

This a project that uses transformers(based on huggingface transformers) to do multi-modal multi-task learning.

Why do I need this

In the original BERT code, neither multi-task learning or multiple GPU training is possible. Plus, the original purpose of this project is NER which dose not have a working script in the original BERT code.

To sum up, compared to the original bert repo, this repo has the following features:

  1. Multimodal multi-task learning(major reason of re-writing the majority of code).
  2. Multiple GPU training
  3. Support sequence labeling (for example, NER) and Encoder-Decoder Seq2Seq(with transformer decoder).

What type of problems are supported?

  • Masked LM and next sentence prediction Pre-train(pretrain)
  • Classification(cls)
  • Sequence Labeling(seq_tag)
  • Multi-Label Classification(multi_cls)
  • Multi-modal Mask LM(mask_lm)

How to run pre-defined problems

There are two types of chaining operations can be used to chain problems.

  • &. If two problems have the same inputs, they can be chained using &. Problems chained by & will be trained at the same time.
  • |. If two problems don't have the same inputs, they need to be chained using |. Problems chained by | will be sampled to train at every instance.

For example, cws|NER|weibo_ner&weibo_cws, one problem will be sampled at each turn, say weibo_ner&weibo_cws, then weibo_ner and weibo_cws will trained for this turn together. Therefore, in a particular batch, some tasks might not be sampled, and their loss could be 0 in this batch.

Please see the examples in notebooks for more details about training, evaluation and export models.




pip install bert-multitask-learning


这是利用transformer(基于huggingface transformers)进行多模态多任务学习的项目.


在原始的BERT代码中, 是没有办法直接用多GPU进行多任务学习的. 另外, BERT并没有给出序列标注和Seq2seq的训练代码.

因此, 和原来的BERT相比, 这个项目具有以下特点:

  1. 多任务学习
  2. 多GPU训练
  3. 序列标注以及Encoder-decoder seq2seq的支持(用transformer decoder)


  • Masked LM和next sentence prediction预训练(pretrain)
  • 单标签分类(cls)
  • 序列标注(seq_tag)
  • 多标签分类(multi_cls)
  • 多模态Mask LM(mask_lm)



  • &. 如果两个任务有相同的输入, 不同标签的话, 那么他们可以&来连接. 被&连接起来的任务会被同时训练.
  • |. 如果两个任务为不同的输入, 那么他们必须|来连接. 被|连接起来的任务会被随机抽取来训练.

例如, 我们定义任务cws|NER|weibo_ner&weibo_cws, 那么在生成每一条数据时, 一个任务块会被随机抽取出来, 例如在这一次抽样中, weibo_ner&weibo_cws被选中. 那么这次weibo_nerweibo_cws会被同时训练. 因此, 在一个batch中, 有可能某些任务没有被抽中, loss为0.

训练, eval和导出模型请见notebooks

Project details

Release history Release notifications | RSS feed

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for bert-multitask-learning, version 0.7.0
Filename, size File type Python version Upload date Hashes
Filename, size bert_multitask_learning-0.7.0-py3-none-any.whl (101.5 kB) File type Wheel Python version py3 Upload date Hashes View
Filename, size bert_multitask_learning-0.7.0.tar.gz (46.6 kB) File type Source Python version None Upload date Hashes View

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page