从头训练大模型,需要海量的数据。但是很多通用的规律其实已经被别人用大模型学好了。所以我们只需要借用一个已经会看世界的模型的眼睛就可以了。
别人做的预训练就是在大数据上学了通用特征,也就是权重,而我们要做的,就是用小数据来学任务特化,在别人的权重上继续训练。这是有现实意义的,比如在生信领域,有标注的基因组数据是稀少的。而无标注的基因组序列是非常多的。
这里用一个预训练模型Nucleotide Transformer v2 (NT v2, 50M, multi-species),用的是ESM架构。这是一个DNA的语言模型,它把DNA序列当成一种语言,每个token是一个6-mer。
| | |
|---|
| 4107 | 4⁶=4096 种 6-mer + 6 个特殊 token(pad/mask/cls 等)+ 5个单碱基(ATCGN) |
| | |
| | |
| | |
| | |
| rotary(RoPE) | |
| | |
| | |
它做的训练就是喂了海量多物种基因组,训练任务是:随机遮住一些6-mer token,让模型根据上下文把它猜出来。模型为了猜对,就被迫学会了DNA的“语法”:motif、GC含量等。这就是它会看基因组的原因。
但它的“头”是预测下一个token,那么这里我们要做就是扔掉这个语言模型头,换成我们要做的任务头。这里我们还是做回 《识别DNA启动子》的任务。那么就要把头换成随机初始化的2分类头,同时保留12层骨干的预训练权重。所以“微调”就是借用它已经学好的DNA通用特征,并教会它这个二分类任务。
加载必要的库
import randomimport torchfrom torch.utils.data import DataLoader, TensorDatasetfrom transformers import AutoTokenizer, AutoModelForSequenceClassification
数据
这个和前面的《用CNN识别DNA启动子》是一样的:
def generate_data(n=2000, length=100, motif="TATAAA"): random.seed(42) pos, neg = [], [] bases = "ACGT" for _ in range(n // 2): seq = "".join(random.choice(bases) for _ in range(length)) i = random.randint(0, length - len(motif)) pos.append(seq[:i] + motif + seq[i + len(motif):]) seq = "".join(random.choice(bases) for _ in range(length)) while motif in seq: seq="".join(random.choice(bases)for_inrange(length)) neg.append(seq) seqs = pos + neg y = torch.tensor([1]*len(pos) + [0]*len(neg), dtype=torch.long) perm = torch.randperm(len(y)) return [seqs[i] for i in perm], y[perm]
tokenizer
def load_tokenizer(): tokenizer = AutoTokenizer.from_pretrained( "InstaDeepAI/nucleotide-transformer-v2-50M-multi-species", trust_remote_code=True) return tokenizer
def tokenize(seqs, tokenizer): enc = tokenizer(seqs, padding=True, truncation=True, max_length=128,return_tensors="pt") return enc["input_ids"], enc["attention_mask"]
首先是加载一个预训练的模型,首次加载会自动去下载。AutoTokenizer我们可以认为是一个泛型类,它会根据加载的模型,去加载相应的类,这里的类是EsmTokenizer。
那么我们可以试一下:
tokenizer = load_tokenizer()tokenize(['TATAAACGTT', 'ACGTACGTACGTACGT'], tokenizer)
它的输出是:
(tensor([[ 3, 1094, 4104, 4105, 4103, 4103, 1], [3,728,3379,4102,4104,4105, 4103]]), tensor([[1, 1, 1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 1, 1]]))
第一个输出就是把序列变成数字ID,按照词表来转换。这个词表是预训练时定死的。我们必须要用它的tokenizer,不然模型的权重就没有用了。第二个输出是告诉我们哪个位置的内容是真内容(1),哪个是 padding(0)。所以后面这个输出也是重要的,不能丢,不然模型会把 当成是真的碱基去算attention,padding位置就混进了结果。
这里是以6-mer为一个token,像第一个序列‘TATAAACGTT’,10个碱基,就是前面6个是一个token,后面4个凑不成6-mer,就拆成单碱基,以这样的规则来做的。
位置: 012345 6TATAAACGTT: =3 TATAAA=1094 C=4104 G=4105 T=4103 T=4103 =1ACGTACGT...: =3 ACGTAC=728 GTACGT A=4102 C=4104 G=4105 T=4103
前面的代表是分类,是整条序列的代言人。
加载预训练模型
def load_model(device): model = AutoModelForSequenceClassification.from_pretrained( "InstaDeepAI/nucleotide-transformer-v2-50M-multi-species", num_labels=2, trust_remote_code=True).to(device)
return(model)
这里的AutoModelForSequenceClassification 会根据模型的名字去找相应的config.json,然后就能确定是该实例化哪个具体的类。在这个具体的例子中,会被解析成EsmForSequenceClassification,我们用真正干活的这个类,也是可以的。但是代码就被写死了,换一个模型换一个类。用这种写法,只要换模型名字就行,具体干活的类根本不用管。这里num_labels=2就触发了换头,它丢掉了语言模型头,新加一个分类头classifier,输出2个logits。
Auto家族:
| | | |
|---|
AutoModel | | | |
AutoModelForSequenceClassification | | | |
AutoModelForMaskedLM | | | |
AutoModelForTokenClassification | | | |
根据不同的任务,我们可以用不同的Auto类。
训练和评估
def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss = 0.0 for batch in loader: input_ids,mask,y=[t.to(device)fortin batch] optimizer.zero_grad() outputs = model(input_ids=input_ids, attention_mask=mask) loss = criterion(outputs.logits, y) loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(loader)
def evaluate(model, loader, device): model.eval() correct, total = 0, 0 with torch.no_grad(): for batch in loader: input_ids,mask,y=[t.to(device)fortin batch] outputs = model(input_ids=input_ids, attention_mask=mask) correct+=(outputs.logits.argmax(dim=1)==y).sum().item() total += y.size(0) return correct / total * 100
这个是没区别的。
主函数
def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"使用设备: {device}")
print("生成数据...") seqs, y = generate_data(8000)
print("加载 tokenizer 与预训练模型(首次下载约 500MB)...") tokenizer = load_tokenizer() model = load_model(device)
print("分词...") input_ids, attention_mask = tokenize(seqs, tokenizer) dataset = TensorDataset(input_ids, attention_mask, y) train_loader = DataLoader(TensorDataset(*[t[:1600] for t
in dataset.tensors]), batch_size=16, shuffle=True) test_loader = DataLoader(TensorDataset(*[t[1600:] for t in dataset.tensors]), batch_size=32)
criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
print("微调 3 个 epoch(预训练模型收敛很快)...") for epoch in range(1, 11): loss = train_one_epoch(model, train_loader, optimizer, criterion, device) acc = evaluate(model, test_loader, device) print(f"Epoch {epoch:2d}/10 | 损失: {loss:.4f} | 准确率: {acc:.2f}%")
print(f"\n完成!准确率 {acc:.2f}%(数据量足够时迁移学习也能到 95%+)")
这些基本上都没有变化,我们省了设计模型这一步,直接用了预训练的。然后其它都基本没区别。训的时候,从预训练的权重开始,而不是随机权重开始,所以收敛会很快,我们的learning rate也要相应地设小一些。因为预训练的权重已经很好了,lr太大会【灾难性遗忘】,把通用特征冲掉。所以微调的lr比从头训练小1-2个数量级。
使用设备: cuda生成数据...加载 tokenizer 与预训练模型(首次下载约 500MB)...
Some weightsofthemodelcheckpointatInstaDeepAI/nucleotide-transformer-v2-50M-multi-specieswerenotusedwheninitializingEsmForSequenceClassification:['lm_head.bias','lm_head.decoder.weight','lm_head.dense.bias','lm_head.dense.weight','lm_head.layer_norm.bias','lm_head.layer_norm.weight']- ThisISexpectedifyouareinitializingEsmForSequenceClassificationfromthecheckpointofamodeltrainedonanothertaskorwithanotherarchitecture(e.g.initializingaBertForSequenceClassificationmodelfromaBertForPreTraining model).- ThisISNOTexpectedifyouareinitializing
EsmForSequenceClassificationfromthecheckpointofamodelthatyouexpecttobeexactlyidentical(initializingaBertForSequenceClassificationmodelfromaBertForSequenceClassification model).Some weightsofEsmForSequenceClassificationwerenotinitializedfromthemodelcheckpointatInstaDeepAI/nucleotide-transformer-v2-50M-multi-speciesandarenewlyinitialized:['classifier.dense.bias','classifier.dense.weight','classifier.out_proj.bias','classifier.out_proj.weight']You shouldprobablyTRAINthismodelonadown-streamtasktobeabletouseitforpredictionsand inference.
分词...微调 3 个 epoch(预训练模型收敛很快)...Epoch 1/10 | 损失: 0.6308 | 准确率: 78.56%Epoch 2/10 | 损失: 0.2605 | 准确率: 93.92%Epoch 3/10 | 损失: 0.0790 | 准确率: 95.22%Epoch 4/10 | 损失: 0.0156 | 准确率: 95.69%Epoch 5/10 | 损失: 0.0067 | 准确率: 95.39%Epoch 6/10 | 损失: 0.0279 | 准确率: 93.78%Epoch 7/10 | 损失: 0.0234 | 准确率: 95.66%Epoch 8/10 | 损失: 0.0100 | 准确率: 96.00%Epoch 9/10 | 损失: 0.0196 | 准确率: 89.23%Epoch 10/10 | 损失: 0.0586 | 准确率: 93.73%
完成!准确率 93.73%(数据量足够时迁移学习也能到 95%+)
这就是输出,中间的message是由于换头所导致的,可以不用管。从这个输出来看,3-4个Epoch就够了,后面明显过拟合了。然后我们还看到epoch 9的96->89这种剧烈抖动的,这个我们不能完全归结于过拟合,还需要有lr scheduler才好。