<think>我们保存模型时使用了`model.save_pretrained`,这应该保存了整个`LlamaWithEncoder`模型的权重和配置。但是,在加载时,我们使用`AutoModelForCausalLM.from_pretrained`只能加载基础模型(即原始的Llama模型部分),而自定义的编码器和融合层不会被自动加载。因此,我们需要修改加载方式,确保整个自定义模型被正确加载。
解决方案:
1. 在保存模型时,我们同时保存了自定义模型的配置(在`config.json`中)和权重(在`safetensors`文件中)。但是,`AutoModelForCausalLM`并不知道我们的自定义类`LlamaWithEncoder`。
2. 我们需要使用`from_pretrained`方法加载整个自定义模型,而不是仅加载基础模型。这要求我们在加载模型之前,将自定义模型类注册到transformers库中。
修改步骤:
- 在推理脚本中,我们需要将`LlamaWithEncoder`类注册为AutoModel的子类,这样`AutoModel`才能识别并加载它。
- 使用`AutoModel`而不是`AutoModelForCausalLM`来加载整个模型。
具体步骤:
1. 在定义`LlamaWithEncoder`类之后,使用`transformers.AutoConfig.register`注册自定义模型。
2. 然后使用`transformers.AutoModel.from_pretrained`加载模型。
但是,由于我们的模型是用于因果语言建模的,我们也可以使用`AutoModelForCausalLM`,但同样需要注册。不过,更简单的方式是直接使用我们自定义的类来加载模型。
另一种方法:我们直接使用`LlamaWithEncoder`类的`from_pretrained`方法(继承自`PreTrainedModel`)来加载模型。但是,我们的`LlamaWithEncoder`类在初始化时需要基础模型,因此不能直接使用`from_pretrained`。
因此,我们修改推理脚本中的模型加载方式:
步骤:
1. 加载基础模型(即原始的Llama模型)作为`base_model`。
2. 然后使用`LlamaWithEncoder`类包裹`base_model`,这样我们就得到了完整的模型结构。
3. 然后,我们需要加载自定义部分的权重(即两个编码器和融合层的权重)。但是,我们保存整个模型时,这些权重已经包含在保存的文件中。因此,我们需要将整个保存的权重加载到我们构建的完整模型中。
但是,我们保存的模型权重是包括基础模型和自定义部分的。所以,我们可以这样:
- 先构建一个空的`LlamaWithEncoder`模型(用基础模型初始化)
- 然后,用保存的权重文件加载整个模型(包括基础模型和自定义部分)
但是,这里有一个问题:我们保存的模型权重文件包含了基础模型的权重和自定义部分的权重。而我们在构建`LlamaWithEncoder`时已经包含了基础模型(通过传入的`base_model`)。所以,我们需要覆盖基础模型的权重和自定义部分的权重。
因此,我们可以这样做:
```python
# 首先加载基础模型(但不加载权重,因为后面整个模型权重会覆盖)
base_model = AutoModelForCausalLM.from_config(config) # 从配置初始化,不加载权重
model = LlamaWithEncoder(base_model)
# 然后加载整个模型的权重
model.load_state_dict(torch.load(os.path.join(args.model_path, "pytorch_model.bin"))) # 或者 safetensors 文件
```
但是,更推荐使用transformers提供的`from_pretrained`方法,但需要确保我们的模型类继承自`PreTrainedModel`并且已经正确配置。
实际上,我们在保存模型时调用了`model.save_pretrained`,所以应该包含一个配置文件,其中指定了模型类型。但是,transformers库默认不知道我们的自定义类。
因此,我们需要在加载模型时,让transformers知道我们的自定义模型类。我们可以使用`AutoModel`并注册自定义类。
注册自定义模型类的方法:
```python
from transformers import AutoConfig, AutoModel
# 注册自定义模型类
AutoConfig.register("llama_with_encoder", LlamaWithEncoderConfig) # 如果我们有自定义配置类,但这里没有
# 但是,我们并没有自定义配置类,我们使用的是基础模型的配置。所以,我们可以使用基础模型的配置,并扩展。
# 实际上,我们在保存模型时,保存的配置是基础模型的配置。因此,我们需要告诉transformers,当加载这个配置时,使用我们的自定义模型类。
# 我们可以这样做:在保存模型时,修改配置,添加一个属性`architectures`为["LlamaWithEncoder"]。这样,加载时就会尝试找这个类。
# 但是,我们也可以手动指定模型类来加载:
model = LlamaWithEncoder.from_pretrained(args.model_path, base_model=None) # 但是这样不行,因为我们的初始化需要base_model参数。
# 所以,我们采用另一种方法:先加载基础模型,再构建自定义模型,然后加载整个模型的权重。
具体步骤:
1. 加载基础模型的配置。
2. 用这个配置初始化基础模型(不加载权重)。
3. 用这个基础模型构建自定义模型(LlamaWithEncoder)。
4. 然后加载整个模型的权重(包括基础模型和自定义部分)到自定义模型中。
代码示例:
```python
from transformers import AutoConfig
# 加载配置
config = AutoConfig.from_pretrained(args.model_path)
# 初始化基础模型(不加载权重)
base_model = AutoModelForCausalLM.from_config(config)
# 创建自定义模型
model = LlamaWithEncoder(base_model)
# 加载整个模型的权重
state_dict = torch.load(os.path.join(args.model_path, "pytorch_model.bin")) # 或者使用safetensors
model.load_state_dict(state_dict)
```
但是,我们保存的模型可能是safetensors格式,所以使用`safe_open`来加载。
但是,transformers的`from_pretrained`方法已经处理了这些,我们更希望使用它。因此,我们需要修改自定义模型类,使其能够被`from_pretrained`直接加载。
为了达到这个目的,我们需要:
1. 自定义模型类继承`PreTrainedModel`。
2. 实现`_init_weights`方法(可选,因为我们的子模块已经初始化)。
3. 提供一个配置类(可选,我们使用基础模型的配置)。
但是,我们的自定义模型类已经继承自`PreTrainedModel`,所以满足条件1。
但是,我们的模型在初始化时需要基础模型,而`from_pretrained`方法在加载时不会传递这个参数。
因此,我们需要修改`LlamaWithEncoder`类的初始化方法,使其能够在不传递基础模型的情况下构建。我们可以通过从配置中读取基础模型类型,然后创建基础模型。
修改`LlamaWithEncoder`的`__init__`方法:
原方法:
```python
def __init__(self, base_model, encoder1_dim=18, encoder2_dim=256, hidden_dim=512):
...
```
改为:
```python
def __init__(self, config, encoder1_dim=18, encoder2_dim=256, hidden_dim=512):
super().__init__(config)
# 从配置初始化基础模型
self.base_model = AutoModelForCausalLM.from_config(config)
# 然后构建自定义部分
...
```
这样,我们就可以使用`LlamaWithEncoder.from_pretrained`来加载整个模型了。
但是,由于我们的模型在保存时包含了基础模型和自定义部分的权重,所以这样修改后,基础模型也会被包含在`LlamaWithEncoder`中,然后整个模型权重一次性加载。
然而,这个改动会改变模型的结构,需要重新训练吗?不需要,因为权重名称应该匹配。但是,我们之前训练保存的权重中的基础模型部分叫做`model`,而现在我们叫做`base_model`,这会导致权重名称不匹配。
因此,我们需要保持模型内部的基础模型变量名为`model`(与之前训练保存的权重中的名称一致)。所以,我们修改:
```python
def __init__(self, config, encoder1_dim=18, encoder2_dim=256, hidden_dim=512):
super().__init__(config)
# 从配置初始化基础模型
self.model = AutoModelForCausalLM.from_config(config)
# 然后构建自定义部分
... # 后续代码保持不变,使用self.model
```
这样,权重名称中的`model.xxx`就可以对应上了。
但是,我们之前训练时,`LlamaWithEncoder`类中的基础模型变量名就是`model`,所以这样修改后,权重名称可以匹配。
因此,我们需要修改自定义模型类的初始化方法,使其只接受一个参数`config`(以及一些可选参数),并在内部构建基础模型。
修改后的`LlamaWithEncoder`类(用于推理)如下:
```python
class LlamaWithEncoder(PreTrainedModel):
def __init__(self, config, encoder1_dim=18, encoder2_dim=256, hidden_dim=512):
super().__init__(config)
# 从配置初始化基础模型
self.model = AutoModelForCausalLM.from_config(config)
# 第一个Transformer Encoder
encoder1_layer = nn.TransformerEncoderLayer(
d_model=encoder1_dim,
nhead=3,
dim_feedforward=hidden_dim,
batch_first=True
)
self.encoder1 = nn.TransformerEncoder(encoder1_layer, num_layers=2)
# 第二个Transformer Encoder
encoder2_layer = nn.TransformerEncoderLayer(
d_model=encoder2_dim,
nhead=8,
dim_feedforward=hidden_dim,
batch_first=True
)
self.encoder2 = nn.TransformerEncoder(encoder2_layer, num_layers=2)
# 投影层
self.proj1 = nn.Linear(encoder1_dim, config.hidden_size)
self.proj2 = nn.Linear(encoder2_dim, config.hidden_size)
# 融合层
self.fusion = nn.Linear(2 * config.hidden_size, config.hidden_size)
# 添加embedding层引用
self.embed_tokens = self.model.get_input_embeddings()
```
这样修改后,我们就可以使用`LlamaWithEncoder.from_pretrained`来加载整个模型了。
但是,注意:我们之前训练保存的模型权重中,基础模型部分的权重名称是以`model.`开头的,而在这个新初始化中,基础模型被命名为`self.model`,所以名称匹配。
因此,在推理脚本中,我们可以这样加载模型:
```python
model = LlamaWithEncoder.from_pretrained(
args.model_path,
torch_dtype=torch.bfloat16,
trust_remote_code=True
)
```
但是,由于我们修改了模型类的初始化方式,训练时保存的配置可能没有指定模型架构为`LlamaWithEncoder`,所以我们需要确保在保存模型时,配置中包含了`architectures`字段为`["LlamaWithEncoder"]`。我们在训练代码保存模型时,可以修改配置:
在训练代码中,保存模型之前:
```python
model.config.architectures = ["LlamaWithEncoder"]
model.save_pretrained(save_directory, safe_serialization=True)
```
但是,如果之前已经保存了模型,那么我们需要在模型目录下的`config.json`中手动添加:
```json
"architectures": [
"LlamaWithEncoder"
],
```
这样,当我们使用`LlamaWithEncoder.from_pretrained`时,transformers库就知道使用哪个类来加载模型。
因此,我们修改训练代码(在保存模型之前):
```python
# 设置模型架构标识
model.config.architectures = ["LlamaWithEncoder"]
# 保存整个模型(包括自定义编码器和融合层)为safetensors格式
model.save_pretrained(save_directory, safe_serialization=True)
```
然后,在推理脚本中,我们就可以直接使用:
```python
model = LlamaWithEncoder.from_pretrained(
args.model_path,
torch_dtype=torch.bfloat16,
trust_remote_code=True
)
```
注意:在推理脚本中,必须确保`LlamaWithEncoder`类已经定义,并且与训练时的结构一致。
因此,总结推理脚本的修改:
1. 修改`LlamaWithEncoder`类的初始化方法,使其只接受`config`,并在内部创建基础模型。
2. 在训练代码保存模型时,设置`model.config.architectures = ["LlamaWithEncoder"]`。
3. 在推理脚本中,使用`LlamaWithEncoder.from_pretrained`加载模型。
但是,如果训练时保存的模型没有设置`architectures`,我们可以在推理脚本中手动指定模型类:
```python
model = LlamaWithEncoder.from_pretrained(args.model_path, ...)
```
只要在推理脚本中定义了同名的类,就可以加载。
由于我们修改了训练代码,所以重新保存模型。如果不想重新训练,可以手动修改保存目录下的`config.json`,添加:
```json
"architectures": [
"LlamaWithEncoder"
],
```
现在,我们修改推理脚本,使用新的加载方式。
注意:在推理脚本中,我们不需要再单独加载基础模型,因为整个模型已经包含在`LlamaWithEncoder`中。
因此,推理脚本中的模型加载部分改为:
```python
# 加载整个自定义模型
model = LlamaWithEncoder.from_pretrained(
args.model_path,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
)
model.eval()
model = model.to('cuda' if torch.cuda.is_available() else 'cpu')
```
但是,我们的自定义模型类`LlamaWithEncoder`在初始化时还有一些额外的参数(如`encoder1_dim`等),我们可以通过配置来传递,或者在加载时传递。但更简单的是在保存模型时,将这些参数保存在配置中。
我们在训练代码中,保存模型之前,将这些参数加入配置:
```python
# 在训练代码中
model.config.encoder1_dim = 18
model.config.encoder2_dim = 256
model.config.hidden_dim = 512
```
然后,在`LlamaWithEncoder`的初始化中,从配置读取:
```python
class LlamaWithEncoder(PreTrainedModel):
def __init__(self, config):
super().__init__(config)
# 从配置中读取自定义参数
encoder1_dim = getattr(config, "encoder1_dim", 18)
encoder2_dim = getattr(config, "encoder2_dim", 256)
hidden_dim = getattr(config, "hidden_dim", 512)
# 从配置初始化基础模型
self.model = AutoModelForCausalLM.from_config(config)
...
```
这样,我们就可以在保存的配置中保留这些参数。
因此,我们修改训练代码,在保存模型前设置这些配置属性:
```python
# 设置自定义参数
model.config.encoder1_dim = 18
model.config.encoder2_dim = 256
model.config.hidden_dim = 512
model.config.architectures = ["LlamaWithEncoder"]
```
然后,在推理脚本中,我们就不需要额外传递这些参数了。
综上所述,我们修改训练代码(保存模型部分)和推理脚本(模型加载部分)以及自定义模型类的初始化方法。
由于时间关系,这里给出推理脚本的完整代码(使用新的模型加载方式):</think># 模型加载问题分析与解决方案
您发现了一个关键问题:在推理代码中,我们使用`AutoModelForCausalLM`加载模型,这只能加载基础Llama模型,而无法加载我们自定义的`LlamaWithEncoder`结构。让我提供完整的解决方案:
## 正确的模型加载方法
我们需要修改模型加载方式,确保整个自定义模型结构被正确加载:
```python
import torch
from transformers import AutoConfig, AutoModel
from safetensors import safe_open
from safetensors.torch import load_file
# 修改后的模型加载函数
def load_custom_model(model_path: str):
# 1. 加载模型配置
config = AutoConfig.from_pretrained(model_path)
# 2. 加载基础模型
base_model = AutoModelForCausalLM.from_pretrained(
model_path,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
config=config
)
# 3. 创建自定义模型实例
model = LlamaWithEncoder(base_model)
# 4. 加载自定义权重(safetensors格式)
model_file = f"{model_path}/model.safetensors"
state_dict = {}
with safe_open(model_file, framework="pt") as f:
for key in f.keys():
state_dict[key] = f.get_tensor(key)
# 5. 加载权重到自定义模型
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
# 检查加载结果
if missing_keys:
print(f"Warning: Missing keys: {missing_keys}")
if unexpected_keys:
print(f"Warning: Unexpected keys: {unexpected_keys}")
return model
# 在main函数中使用
def main():
# ... [前面的代码不变] ...
# 加载自定义模型
model = load_custom_model(args.model_path)
model.eval()
model = model.to('cuda' if torch.cuda.is_available() else 'cpu')
# ... [后面的代码不变] ...
```
## 完整推理代码(修正版)
```python
import json
import torch
import pandas as pd
import re
import os
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
from safetensors import safe_open
import argparse
import numpy as np
import torch.nn as nn
from tqdm import tqdm
# 常量定义 [与之前相同]...
# 化学式转密集向量 [与之前相同]...
# 位置编码生成 [与之前相同]...
# 初始化位置编码矩阵
P = positional_encoding(2000000, 254)
dimn = 254 # 与位置编码维度一致
# 质谱数据编码 (单个样本) [与之前相同]...
# 质谱数据预处理 (单个样本) [与之前相同]...
# 定义模型结构 (必须与训练时相同)
class LlamaWithEncoder(nn.Module): # 修改为继承nn.Module
def __init__(self, base_model, encoder1_dim=18, encoder2_dim=256, hidden_dim=512):
super().__init__()
self.config = base_model.config
self.model = base_model
# 第一个Transformer Encoder
encoder1_layer = nn.TransformerEncoderLayer(
d_model=encoder1_dim,
nhead=3,
dim_feedforward=hidden_dim,
batch_first=True
)
self.encoder1 = nn.TransformerEncoder(encoder1_layer, num_layers=2)
# 第二个Transformer Encoder
encoder2_layer = nn.TransformerEncoderLayer(
d_model=encoder2_dim,
nhead=8,
dim_feedforward=hidden_dim,
batch_first=True
)
self.encoder2 = nn.TransformerEncoder(encoder2_layer, num_layers=2)
# 投影层
self.proj1 = nn.Linear(encoder1_dim, base_model.config.hidden_size)
self.proj2 = nn.Linear(encoder2_dim, base_model.config.hidden_size)
# 融合层
self.fusion = nn.Linear(2 * base_model.config.hidden_size, base_model.config.hidden_size)
# 添加embedding层引用
self.embed_tokens = base_model.get_input_embeddings()
# 添加生成方法
def generate(self, input_ids, encoder1_inputs, encoder2_inputs, **kwargs):
# 准备输入嵌入
embeddings = self.embed_tokens(input_ids)
# 处理编码器输入
enc1_out = self.encoder1(encoder1_inputs)
enc1_out = enc1_out.mean(dim=1)
enc1_proj = self.proj1(enc1_out)
enc2_out = self.encoder2(encoder2_inputs)
enc2_out = enc2_out.mean(dim=1)
enc2_proj = self.proj2(enc2_out)
# 融合编码器输出
fused = self.fusion(torch.cat([enc1_proj, enc2_proj], dim=1))
fused = fused.unsqueeze(1)
# 将融合结果与第一个token的嵌入结合
if embeddings.size(1) > 0:
embeddings[:, 0, :] = 0.7 * embeddings[:, 0, :] + 0.3 * fused[:, 0, :]
# 调用基础模型的生成方法
return self.model.generate(
inputs_embeds=embeddings,
**kwargs
)
# 前向传播 [与之前相同]...
# 自定义模型加载函数
def load_custom_model(model_path: str):
# 1. 加载模型配置
config = AutoConfig.from_pretrained(model_path)
# 2. 加载基础模型
base_model = AutoModelForCausalLM.from_pretrained(
model_path,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
config=config
)
# 3. 创建自定义模型实例
model = LlamaWithEncoder(base_model)
# 4. 加载自定义权重(safetensors格式)
model_file = f"{model_path}/model.safetensors"
if os.path.exists(model_file):
state_dict = {}
with safe_open(model_file, framework="pt") as f:
for key in f.keys():
state_dict[key] = f.get_tensor(key)
# 5. 加载权重到自定义模型
model.load_state_dict(state_dict, strict=False)
else:
print(f"Warning: No model.safetensors found in {model_path}, using base model weights only")
return model
def main():
# 解析命令行参数
parser = argparse.ArgumentParser(description='Run inference with the trained model')
parser.add_argument('--model_path', type=str, default='./llama3.2-SELFIES',
help='Path to the saved model directory')
parser.add_argument('--input_csv', type=str, required=True,
help='Path to input CSV file for inference')
parser.add_argument('--output_csv', type=str, default='predictions.csv',
help='Path to output CSV file with predictions')
parser.add_argument('--batch_size', type=int, default=8,
help='Batch size for inference')
args = parser.parse_args()
# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained(args.model_path)
# 加载自定义模型
model = load_custom_model(args.model_path)
model.eval()
model = model.to('cuda' if torch.cuda.is_available() else 'cpu')
print("Model loaded successfully")
# 加载数据
df = pd.read_csv(args.input_csv)
# 准备预测结果容器
predictions = []
# 分批处理
for i in tqdm(range(0, len(df), args.batch_size), desc="Processing batches"):
batch_df = df.iloc[i:i+args.batch_size]
batch_data = []
# 预处理每个样本
for idx, row in batch_df.iterrows():
formula = row['Molecular Formula']
formula_vec = formula_to_dense(formula).unsqueeze(0).unsqueeze(0) # (1, 1, 18)
# 预处理质谱
spectrum = preprocess_single_spectrum(row['Spectrum'], row['Total Exact Mass'])
spec_matrix = encode_single_spectrum(spectrum, P, dimn).unsqueeze(0) # (1, 501, 258)
# 创建起始token
input_ids = torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long)
batch_data.append({
'formula_vec': formula_vec,
'spec_matrix': spec_matrix,
'input_ids': input_ids
})
# 准备批次输入
formula_vecs = torch.cat([d['formula_vec'] for d in batch_data], dim=0)
spec_matrices = torch.cat([d['spec_matrix'] for d in batch_data], dim=0)
input_ids = torch.cat([d['input_ids'] for d in batch_data], dim=0)
# 移动到设备
device = next(model.parameters()).device
formula_vecs = formula_vecs.to(device)
spec_matrices = spec_matrices.to(device)
input_ids = input_ids.to(device)
# 生成SELFIES序列
with torch.no_grad():
generated = model.generate(
input_ids=input_ids,
encoder1_inputs=formula_vecs,
encoder2_inputs=spec_matrices,
max_length=512,
num_beams=3,
early_stopping=True,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
)
# 解码生成的序列
batch_preds = tokenizer.batch_decode(generated, skip_special_tokens=True)
predictions.extend(batch_preds)
# 保存结果到CSV
df['Predicted_SELFIES'] = predictions
df.to_csv(args.output_csv, index=False)
print(f"Predictions saved to {args.output_csv}")
if __name__ == "__main__":
main()
```
## 关键修改说明
### 1. 模型加载流程重构
- **自定义加载函数**:创建`load_custom_model`函数处理整个加载过程
- **分步加载**:
1. 加载配置
2. 加载基础Llama模型
3. 创建自定义模型实例
4. 加载自定义权重
- **safetensors支持**:使用`safe_open`安全加载模型权重
### 2. 模型结构修改
- **继承nn.Module**:将`PreTrainedModel`改为继承`nn.Module`,简化模型管理
- **添加generate方法**:在自定义模型中实现生成接口,封装基础模型的生成功能
- **权重加载处理**:使用`load_state_dict`加载保存的自定义权重
### 3. 生成流程优化
- **设备自动检测**:`device = next(model.parameters()).device`获取模型所在设备
- **输入预处理**:确保所有输入张量都在正确设备和数据类型上
- **错误处理**:添加对缺失权重文件的警告
### 4. 模型保存修正
在训练代码中,确保正确保存整个模型:
```python
# 保存整个模型(包括自定义编码器和融合层)为safetensors格式
save_directory = './llama3.2-SELFIES'
# 保存模型权重
torch.save(model.state_dict(), f"{save_directory}/model.safetensors")
# 保存配置
model.config.save_pretrained(save_directory)
# 保存tokenizer
tokenizer.save_pretrained(save_directory)
```
## 使用说明
### 1. 训练后保存模型
```python
# 在训练脚本末尾添加
model.save_pretrained(save_directory) # 保存基础模型
torch.save(model.state_dict(), f"{save_directory}/custom_weights.safetensors") # 保存自定义权重
```
### 2. 运行推理
```bash
python inference_custom.py \
--model_path ./llama3.2-SELFIES \
--input_csv test_data.csv \
--output_csv predictions.csv \
--batch_size 4
```
### 3. 目录结构要求
模型目录应包含:
```
llama3.2-SELFIES/
├── config.json
├── tokenizer.json
├── tokenizer_config.json
├── special_tokens_map.json
├── model.safetensors # 自定义权重
└── pytorch_model.bin # 基础模型权重
```
## 常见问题解决
1. **权重不匹配错误**:
```python
# 在load_custom_model中添加:
model.load_state_dict(state_dict, strict=False)
```
忽略不匹配的键,只加载匹配的部分
2. **设备内存不足**:
- 减小`batch_size`
- 使用`fp16`代替`bf16`:
```python
model = model.half()
```
3. **缺少safetensors文件**:
```python
# 在训练代码中确保保存:
torch.save(model.state_dict(), "model.safetensors")
```