235 lines
7.5 KiB
Python
235 lines
7.5 KiB
Python
import os
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = '2,3'
|
||
|
||
import pickle
|
||
from PIL import Image
|
||
import torch
|
||
from torch.utils.data import Dataset, DataLoader
|
||
from predata import prepare_training_data
|
||
from transformers import AutoTokenizer, AutoModel
|
||
import torch.nn as nn
|
||
from torch.optim import AdamW
|
||
|
||
# 加载tokenizer
|
||
tokenizer = AutoTokenizer.from_pretrained("Internvl2_5")
|
||
|
||
# 加载处理好的数据
|
||
print("正在加载数据...")
|
||
with open('valid_products.pkl', 'rb') as f:
|
||
data = pickle.load(f)
|
||
|
||
print(f"成功加载数据，共 {len(data)} 条记录")
|
||
|
||
class ProductDataset(Dataset):
|
||
def __init__(self, data, tokenizer):
|
||
self.data = data
|
||
self.tokenizer = tokenizer
|
||
|
||
def __len__(self):
|
||
return len(self.data)
|
||
|
||
def __getitem__(self, idx):
|
||
item = self.data[idx]
|
||
question = "Tell me the product name in the picture."
|
||
answer = "The product name is: " + item['name']
|
||
|
||
# 直接使用prepare_training_data处理所有数据
|
||
training_data = prepare_training_data(
|
||
question=question,
|
||
answer=answer,
|
||
tokenizer=self.tokenizer,
|
||
image_path=item['image'] # 直接传入PIL Image对象
|
||
)
|
||
|
||
return training_data
|
||
|
||
# 创建数据集实例
|
||
dataset = ProductDataset(data, tokenizer)
|
||
|
||
# 测试输出第一条数据看看
|
||
sample = dataset[0]
|
||
print("\n数据样例:")
|
||
print(f"输入形状: {sample['input_ids'].shape}")
|
||
print(f"图片张量形状: {sample['pixel_values'].shape}")
|
||
print(f"标签形状: {sample['labels'].shape}")
|
||
|
||
# 设置设备
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
# 加载模型
|
||
path = 'Internvl2_5'
|
||
model = AutoModel.from_pretrained(
|
||
path,
|
||
torch_dtype=torch.bfloat16,
|
||
low_cpu_mem_usage=True,
|
||
use_flash_attn=True,
|
||
trust_remote_code=True,
|
||
vision_model = None,
|
||
language_model = None).train().cuda()
|
||
tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)
|
||
|
||
# 加载之前训练的权重
|
||
print("正在加载预训练权重 model_epoch_5.pth ...")
|
||
model.load_state_dict(torch.load('model_epoch_5.pth'))
|
||
print("成功加载预训练权重")
|
||
|
||
def prepare_model_for_training(model, cast_trainable_params_to_fp32=True):
|
||
print("正在设置模型参数...")
|
||
|
||
# 冻结vision_model
|
||
for param in model.vision_model.parameters():
|
||
param.requires_grad_(False)
|
||
|
||
# 冻结language_model
|
||
for param in model.language_model.parameters():
|
||
param.requires_grad_(False)
|
||
|
||
# 设置mlp1为可训练，并可选转换为fp32
|
||
# for param in model.mlp1.parameters():
|
||
# if cast_trainable_params_to_fp32:
|
||
# param.data = param.data.to(torch.float32)
|
||
|
||
# 打印可训练参数数量
|
||
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
print(f"可训练参数数量: {trainable_params}")
|
||
|
||
return model
|
||
|
||
# 准备模型
|
||
model = prepare_model_for_training(model)
|
||
|
||
# 只对可训练参数创建优化器
|
||
optimizer = AdamW(
|
||
[p for p in model.parameters() if p.requires_grad],
|
||
lr=1e-4, # 学习率
|
||
betas=(0.9, 0.999),
|
||
eps=1e-8,
|
||
weight_decay=0.01
|
||
)
|
||
|
||
def custom_collate_fn(batch):
|
||
"""
|
||
自定义collate函数来处理变长序列
|
||
"""
|
||
# 获取batch中最大的序列长度
|
||
max_len = max([b['input_ids'].size(1) for b in batch])
|
||
|
||
batch_size = len(batch)
|
||
|
||
# 创建填充后的张量
|
||
input_ids = torch.zeros((batch_size, max_len), dtype=batch[0]['input_ids'].dtype)
|
||
attention_mask = torch.zeros((batch_size, max_len), dtype=batch[0]['attention_mask'].dtype)
|
||
labels = torch.full((batch_size, max_len), -100, dtype=batch[0]['labels'].dtype) # 用-100填充标签
|
||
|
||
# 填充每个样本
|
||
for i, item in enumerate(batch):
|
||
seq_len = item['input_ids'].size(1)
|
||
input_ids[i, :seq_len] = item['input_ids'][0, :seq_len]
|
||
attention_mask[i, :seq_len] = item['attention_mask'][0, :seq_len]
|
||
labels[i, :seq_len] = item['labels'][0, :seq_len]
|
||
|
||
# 处理pixel_values (这个应该是固定大小的)
|
||
pixel_values = torch.stack([item['pixel_values'] for item in batch])
|
||
|
||
return {
|
||
'input_ids': input_ids,
|
||
'attention_mask': attention_mask,
|
||
'labels': labels,
|
||
'pixel_values': pixel_values
|
||
}
|
||
|
||
# 使用自定义的collate_fn创建DataLoader
|
||
train_loader = DataLoader(
|
||
dataset,
|
||
batch_size=5, # 可以根据需要调整
|
||
shuffle=True,
|
||
pin_memory=torch.cuda.is_available(),
|
||
collate_fn=custom_collate_fn # 使用自定义的collate函数
|
||
)
|
||
|
||
# 计算img_context_token_id
|
||
IMG_CONTEXT_TOKEN = '<IMG_CONTEXT>'
|
||
img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
|
||
|
||
def extract_feature(model, pixel_values):
|
||
# 移除多余的维度
|
||
if pixel_values.dim() == 5: # [batch_size, 1, 3, 448, 448]
|
||
pixel_values = pixel_values.squeeze(1) # 变成 [batch_size, 3, 448, 448]
|
||
|
||
# 使用vision_model提取特征
|
||
vit_embeds = model.vision_model(
|
||
pixel_values=pixel_values,
|
||
output_hidden_states=False,
|
||
return_dict=True
|
||
).last_hidden_state
|
||
|
||
vit_embeds = vit_embeds[:, 1:, :] # 移除CLS token
|
||
|
||
# 重塑并处理特征
|
||
h = w = int(vit_embeds.shape[1] ** 0.5)
|
||
vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
|
||
vit_embeds = model.pixel_shuffle(vit_embeds, scale_factor=model.downsample_ratio)
|
||
vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
|
||
vit_embeds = model.mlp1(vit_embeds)
|
||
return vit_embeds
|
||
|
||
# 训练循环
|
||
num_epochs = 40 # 修改为40个epochs
|
||
for epoch in range(num_epochs):
|
||
model.train()
|
||
total_loss = 0
|
||
for batch_idx, batch in enumerate(train_loader):
|
||
# 将数据移到GPU
|
||
input_ids = batch['input_ids'].to(device)
|
||
attention_mask = batch['attention_mask'].to(device)
|
||
labels = batch['labels'].to(device)
|
||
pixel_values = batch['pixel_values'].to(device)
|
||
|
||
# 提取图像特征
|
||
vit_embeds = extract_feature(model, pixel_values)
|
||
|
||
# 计算输入嵌入
|
||
input_embeds = model.language_model.get_input_embeddings()(input_ids)
|
||
B, N, C = input_embeds.shape
|
||
input_embeds = input_embeds.reshape(B * N, C)
|
||
|
||
# 替换图像上下文token的嵌入
|
||
input_ids_flat = input_ids.reshape(B * N)
|
||
selected = (input_ids_flat == img_context_token_id)
|
||
input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)
|
||
|
||
# 恢复原始形状
|
||
input_embeds = input_embeds.reshape(B, N, C)
|
||
|
||
# 前向传播
|
||
outputs = model.language_model(
|
||
inputs_embeds=input_embeds,
|
||
attention_mask=attention_mask,
|
||
labels=labels,
|
||
return_dict=True
|
||
)
|
||
|
||
loss = outputs.loss
|
||
total_loss += loss.item()
|
||
|
||
# 反向传播
|
||
loss.backward()
|
||
optimizer.step()
|
||
optimizer.zero_grad()
|
||
|
||
# 每100步打印一次loss
|
||
if batch_idx % 100 == 0:
|
||
print(f'Epoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item():.4f}')
|
||
|
||
# 每个epoch结束打印平均loss
|
||
avg_loss = total_loss / len(train_loader)
|
||
print(f'Epoch {epoch} completed. Average Loss: {avg_loss:.4f}')
|
||
|
||
# 每10个epoch保存一次模型
|
||
if (epoch + 1) % 10 == 0:
|
||
save_path = f'model_epoch_{epoch+1}.pth'
|
||
print(f"保存模型权重到 {save_path}")
|
||
torch.save(model.state_dict(), save_path)
|
||
|
||
|