CogVLM训练源码解读--数据处理

---

前言

本文是CogVLM是一个多模态大型模型，它能够处理文本、图像和其他类型的数据。在数据处理方面，CogVLM可以接收多种类型的输入数据，包括文本、图像、音频等。然而，很少有人对代码数据处理进行解读或者基本找不到。基于此，本文将结合源码给出CogVLM大模型数据处理内容，主要包含图像数据处理、文本tokenizer构建、文本加工与修改自己文本方法代码修改。总之，我将结合代码一步一步带领读者实现大模型数据处理源码内容。

---

一、数据主函数源码解读

CogVLM的数据处理包含2部分，一部分是图像数据处理，另一部分是文本数据处理。其源码位于finetune_cogvlm_demo.py文件，如下：

from utils.utils import llama2_tokenizertokenizer = llama2_tokenizer(args.local_tokenizer, signal_type=args.version)image_processor = get_image_processor(args.eva_args["image_size"][0])  # 获得图像加工函数，并附image_size参数text_processor = llama2_text_processor(tokenizer, args.max_length, args.image_length)  # 获得文本加工函数model = training_main(args, model_cls=model, forward_step_function=forward_step, create_dataset_function=partial(create_dataset_function, image_processor, text_processor), collate_fn=data_collator, forward_step_eval=forward_step_eval)#        训练函数      参数         模型                    训练机制                                    处理数据函数-->给该函数添加部分参数<---参数为函数                               dataloader整合数据                 评估函数               

从上面可看出获得图像处理函数调用是通过get_image_processor,而获得文本处理函数是调用llama2_text_processor函数，然后通过training_main函数把图像处理与文本处理函数分别作为image_processor与text_processor参数传递。我将说明这2个函数如何处理成参数方法。

1、图像函数源码调用解读

图像处理实际是调用blip2_image_processor_func_with_inputs函数(后面我会详细解读)，实现图像加工，最终作为image_processor函数参数。源码如下：

def get_image_processor(image_size):return partial(blip2_image_processor_func_with_inputs, BlipImageEvalProcessor(image_size))

2、文本函数源码调用解读

文本处理实际是通过huggingface的方式调用tokenizer等方法，在通过llama2_text_processor类对文本处理，我后面会解读，这里介绍通过这样方式作为text_processor函数参数。源码如下：

    from utils.utils import llama2_tokenizertokenizer = llama2_tokenizer(args.local_tokenizer, signal_type=args.version)text_processor = llama2_text_processor(tokenizer, args.max_length, args.image_length)  # 获得文本加工函数

3、tokenizer生成函数

对文本的tokenizer处理，CogVLM调用huggingface的llama函数，使用from utils.utils import llama2_tokenizer调用，在源码中使用函数如下：

tokenizer = llama2_tokenizer(args.local_tokenizer, signal_type=args.version) 

而对于语言函数

from transformers import LlamaTokenizer
def llama2_tokenizer(tokenizer_path, signal_type="base"):tokenizer = LlamaTokenizer.from_pretrained(tokenizer_path)if tokenizer.pad_token_id is None:tokenizer.pad_token_id = 32000tokenizer.boi = "[IMG]"tokenizer.eoi = "[/IMG]"assert signal_type in ["base", "chat", "vqa", "chat_old"]tokenizer.signal_type = signal_typereturn tokenizer

4、llama2_text_processor文本处理函数解读

该函数是文本处理相关内容，我将其代码罗列如下：

class llama2_text_processor:def __init__(self, tokenizer, max_target_length=2048, image_length=257, model=None):self.tokenizer = tokenizerself.max_target_length = max_target_lengthself.image_length = image_lengthdef __call__(self, caption, prompt=""):if '<EOI>' not in prompt:prompt = self.replace_tags_with_empty(prompt)# caption = self.replace_tags_with_empty(caption)history = []prompt = self.history_to_prompt(prompt, history)input_ids = [self.tokenizer.bos_token_id]prompt_splits = prompt.split('<EOI>')caption_splits = caption.split('<EOI>')if len(prompt_splits) > 0:input_ids.extend(self.tokenizer.encode(prompt_splits[0], add_special_tokens=False))for tokens in prompt_splits[1:]:tokens_with_img = [-100] + self.tokenizer.encode(tokens, add_special_tokens=False)input_ids.extend(tokens_with_img)context_length = len(input_ids) + (len(prompt_splits)-1) * (self.image_length + 1)if context_length > self.max_target_length - 10:return Noneif len(caption_splits) > 0:input_ids.extend(self.tokenizer.encode(caption_splits[0], add_special_tokens=False))for tokens in caption_splits[1:]:tokens_with_img = [-100] + self.tokenizer.encode(tokens, add_special_tokens=False)input_ids.extend(tokens_with_img)if len(input_ids) > self.max_target_length - self.image_length - 5:input_ids = input_ids[:self.max_target_length - self.image_length - 5]input_ids += [self.tokenizer.eos_token_id]while -100 in input_ids:img_idx = input_ids.index(-100)input_ids = input_ids[:img_idx] + [0] * (self.image_length + 1) + [-1] + input_ids[img_idx+1:]image_position = []while -1 in input_ids:img_idx = input_ids.index(-1)input_ids[img_idx] = 0image_position.append(img_idx)image_embed_mask = [0] * len(input_ids)vision_expert_mask = [0] * len(input_ids)image_rope_mask = [0] * len(input_ids)for idx in image_position:image_embed_mask[idx-self.image_length-1: idx+1] = [1] * (self.image_length + 2)vision_expert_mask[idx-self.image_length-1: idx] = [1] * (self.image_length + 1)image_rope_mask[idx - self.image_length: idx] = [1] * self.image_lengthattention_mask = [1] * len(input_ids)labels = [-100] * context_length + input_ids[context_length:]pad_len = self.max_target_length - len(input_ids)input_ids = input_ids + [self.tokenizer.pad_token_id] * pad_lenattention_mask = attention_mask + [1] * pad_lenvision_expert_mask = vision_expert_mask + [0] * pad_lenimage_embed_mask = image_embed_mask + [0] * pad_lenimage_rope_mask = image_rope_mask + [0] * pad_lennp_mask = np.tril(np.expand_dims(np.array(attention_mask), 0).repeat(len(attention_mask), 0))labels = labels + [-100] * pad_lenfor idx in image_position:labels[idx-self.image_length-1: idx+1] = [-100] * (self.image_length + 2)position_ids = []pid = -1for i in range(len(input_ids)):if image_rope_mask[i] == 0 or (i > 0 and image_rope_mask[i] != image_rope_mask[i - 1]):pid += 1position_ids.append(pid)input_ids = torch.tensor(input_ids).unsqueeze(0)labels = torch.tensor(labels).unsqueeze(0)attention_mask = torch.from_numpy(np_mask).unsqueeze(0).unsqueeze(0)image_embed_mask = torch.tensor(image_embed_mask).unsqueeze(0)vision_expert_mask = torch.tensor(vision_expert_mask).unsqueeze(0)image_rope_mask = torch.tensor(image_rope_mask).unsqueeze(0)position_ids = torch.tensor(position_ids).unsqueeze(0)context_length = torch.tensor(context_length).unsqueeze(0).long()return {'input_ids': input_ids, 'labels': labels, 'position_ids': position_ids, 'attention_mask': attention_mask, 'image_embed_mask': image_embed_mask,'context_length': context_length, 'image_position': image_position, 'vision_expert_mask': vision_expert_mask, 'image_rope_mask': image_rope_mask}def history_to_prompt(self, query, history):return _history_to_prompt[self.tokenizer.signal_type](self, query, history)def replace_tags_with_empty(self, text):return re.sub('<pad>|<s>|</s>|<EOI>', '', text)

二、create_dataset_function函数源码代码解读

在使用model = training_main函数create_dataset_function=partial(create_dataset_function, image_processor, text_processor)，借助python自带partial函数实现数据类处理，其源码如下：

from utils.utils import ItemDataset
def create_dataset_function(image_processor, text_processor, path, args):dataset = ItemDataset(image_processor, text_processor, args, path)return dataset

我们可发现该函数调用了ItemDataset类，该类恰好是数据加工的迭代器类，类似taorch的dataset等功能，实际也是继承torch的dataset类，进一步封装加工数据。其中image_processor, text_processor是上面提到图像加工与文本加工方法代码。

三、sat库之make_loaders函数源码解读

在上面，我们已给出training_main函数，该函数集成数据处理、模型训练等方法，其中hooks保存不同方法，而数据处理代码为make_loaders。

1、make_loaders函数调用说明

make_loaders为模型数据加载方法，又是sat库集成方法，其源码如下：

# Data stuff. 数据处理方法适用train_data, val_data, test_data = make_loaders(args, hooks['create_dataset_function'], collate_fn=collate_fn) # 通过该函数调用

很明显，make_loaders参数为需要参数、处理数据类(集成dataset)、collate_fn函数(类似dataloader处理batch数据方式)。

源码位置：sat.data_utils.configure_data.py文件

2、make_loaders函数源码解读

我将maker_loaders源码分为三个部分，第一个部分是使用partial函数将传递数据类create_dataset_function（实际是ItemDataset类）赋参数并重命名函数为make_dataset；第二部分是否传训练、验证、测试路径，使用ItemDataset类处理数据，类似torch的dataset结构；第三部分也是调用sat库的make_data_loader包装对应数据，类似torch的dataloader结构，且调用自己传递的collate_fn方法。

第一部分源码解读

将我们对数据处理create_dataset_function类（实际是ItemDataset类）使用sat库的make_dataset_full函数包装，源码如下：

    make_dataset = partial(make_dataset_full, create_dataset_function=create_dataset_function, batch_from_same_dataset=args.batch_from_same_dataset)

第二部分源码解读

根据传递路径参数，使用上面make_dataset方法处理数据，类似torch的dataset，源码如下：

# make datasets splits and tokenizertrain = Nonevalid = Nonetest = Noneif args.train_data is not None:train = make_dataset(**data_set_args, args=args, dataset_weights=args.train_data_weights, is_train_data=True)if should_split(split):train, valid, test = train# make training and val dataset if necessaryif valid is None and args.valid_data is not None:eval_set_args['path'] = args.valid_datavalid = make_dataset(**eval_set_args, args=args, random_mapping=not args.strict_eval)if test is None and args.test_data is not None:eval_set_args['path'] = args.test_datatest = make_dataset(**eval_set_args, args=args, random_mapping=not args.strict_eval)

第三部分源码解读

将我们处理的datset数据进行dataloader包装，类似torch的dataloader，且调用自己写的collate_fn方法，源码如下：

# wrap datasets with data loaderif train is not None and args.batch_size > 0:train = make_data_loader(train, batch_size, args, split='train', collate_fn=collate_fn)args.do_train = Trueelse:args.do_train = Falseeval_batch_size = eval_batch_size if eval_batch_size != 0 else batch_sizeif valid is not None:valid = make_data_loader(valid, eval_batch_size, args, split='val', collate_fn=collate_fn)args.do_valid = Trueelse:args.do_valid = Falseif test is not None:test = make_data_loader(test, eval_batch_size, args, split='test', collate_fn=collate_fn)args.do_test = Trueelse:args.do_test = False

make_loaders源码展示

另外，该函数也对相应参数做处理，如world_size等内容，所有源码如下：

def make_loaders(args, create_dataset_function, collate_fn=None):"""makes training/val/testArgs:args.train_data, args.valid_data, args.test_data: str. Paths to the dataset.args.split: str. format: "8,1,1". how to split train_data.args.dataset_type: use to create the right datasets. """make_dataset = partial(make_dataset_full, create_dataset_function=create_dataset_function, batch_from_same_dataset=args.batch_from_same_dataset)world_size = torch.distributed.get_world_size(   group=mpu.get_data_parallel_group())batch_size = args.batch_size * world_sizeeval_batch_size = batch_sizeif args.eval_batch_size is not None:eval_batch_size = args.eval_batch_size * world_sizesplit = get_split(args)data_set_args = {'path': args.train_data,'split': split,}eval_set_args = copy.copy(data_set_args)eval_set_args['split'] = [1.]# make datasets splits and tokenizertrain = Nonevalid = Nonetest = Noneif args.train_data is not None:train = make_dataset(**data_set_args, args=args, dataset_weights=args.train_data_weights, is_train_data=True)if should_split(split):train, valid, test = train# make training and val dataset if necessaryif valid is None and args.valid_data is not None:eval_set_args['path'] = args.valid_datavalid = make_dataset(**eval_set_args, args=args, random_mapping=not args.strict_eval)if test is None and args.test_data is not None:eval_set_args['path'] = args.test_datatest = make_dataset(**eval_set_args, args=args, random_mapping=not args.strict_eval)# wrap datasets with data loaderif train is not None and args.batch_size > 0:train = make_data_loader(train, batch_size, args, split='train', collate_fn=collate_fn)args.do_train = Trueelse:args.do_train = Falseeval_batch_size = eval_batch_size if eval_batch_size != 0 else batch_sizeif valid is not None:valid = make_data_loader(valid, eval_batch_size, args, split='val', collate_fn=collate_fn)args.do_valid = Trueelse:args.do_valid = Falseif test is not None:test = make_data_loader(test, eval_batch_size, args, split='test', collate_fn=collate_fn)args.do_test = Trueelse:args.do_test = Falsereturn train, valid, test

四、sat库之make_dataset_full函数源码解读

1、参数配置

我是在vscode运行，我对训练数据参数配置作为列子，可配置三种方式，假设CogVLM-SFT-311K文件夹有2个文件，分别为llava_instruction_multi_conversations_formate与llava_instruction_single_conversation_formate文件，可单独给子文件路径、也可给子文件上一个文件路径、也可给多个文件列表路径，具体如下：

# 第一种方式：
"--train-data", "/extend_disk/tj/data/CogVLM-SFT-311K/llava_instruction_multi_conversations_formate",
# 第二种方式：
"--train-data", "/extend_disk/tj/data/CogVLM-SFT-311K",
# 第三种方式：
"--train-data", "/extend_disk/tj/data/CogVLM-SFT-311K/llava_instruction_multi_conversations_formate", "/extend_disk/tj/data/CogVLM-SFT-311K/llava_instruction_single_conversation_formate" ,           

2、数据格式说明

单个数据文件夹内容如下：

json文件内容如下：

3、make_dataset_full数据读取源码解读

make_loaders函数中的make_dataset类调用是被make_dataset_full函数包装，主要处理一些逻辑，使其进入数据类为统一格式，其源码如下：

ds = []
for p in path:d = create_dataset_function(p, args)ds.append(d)
ds = ConcatDataset(ds, weights=dataset_weights)

以上，可看到path类似给定参数，然后使用os.walk遍历所有.jpg格式数据，并给成绝对路径，且将所有数据cat为一个列表。

4、make_dataset_full源码展示

def make_dataset_full(path, split, args, create_dataset_function, dataset_weights=None, random_mapping=True, is_train_data=False, batch_from_same_dataset=False, **kwargs):"""function to create datasets+tokenizers for common options"""print_all('make dataset ' + str(path), level='DEBUG')assert isinstance(path, list)if args.iterable_dataset: # cannot indexed# the random mapping is flexible and efficient, but sometimes we have pratical issue# For instance, someone just gives you a iterable dataset, e.g. webdatasetfrom .webds import ConfiguredResampledShards, DataPipelinevalid_types = (ConfiguredResampledShards, DataPipeline)assert split[0] == 1, 'Iterable dataset cannot auto split.'ds = []for p in path:d = create_dataset_function(p, args)assert isinstance(d, valid_types)ds.append(d)# ds = ChainDataset(ds) # please merge them in a url if chainif batch_from_same_dataset:assert args.num_workers <= 1, 'We cannot control the actual speed of different workers, may mix different iterable parts.'ds = AlterDataset(ds, weights=dataset_weights, seed=args.seed, batch_from_same_dataset=batch_from_same_dataset, batch_size=args.batch_size)return dsif split is None:split = [1.] if not should_split(split):ds = []for p in path:d = create_dataset_function(p, args)ds.append(d)ds = ConcatDataset(ds, weights=dataset_weights)if random_mapping:if args.epochs is not None: # not auto-scale, but use a given number of epoches.ds = RandomDataset(ds, scale=args.epochs, seed=args.seed)else:world_size = torch.distributed.get_world_size(group=mpu.get_data_parallel_group())if is_train_data:# only train-dataset will set this to True,# so we enlarge it to make sure that the data is sufficient.scale = max(200, 1 + (args.train_iters * args.batch_size * args.gradient_accumulation_steps * world_size) // len(ds))else:scale = max(200, 1 + ((1 + args.train_iters // args.eval_interval) * args.eval_iters * args.eval_batch_size * args.gradient_accumulation_steps * world_size) // len(ds))ds = RandomMappingDataset(ds, scale=scale)return ds   [-1, 9, C3, [512]],else:# must first split datasets, then reweight/concat, finally random-mapping.# this order avoids overlapping.train_ds, valid_ds, test_ds = [], [], []for p in path:d = create_dataset_function(p, args)if should_split(split):dtrain, dvalid, dtest = split_ds(d, split, block_size=args.block_size, seed=args.seed)train_ds.append(dtrain)valid_ds.append(dvalid)test_ds.append(dtest)train_ds = ConcatDataset(train_ds, weights=dataset_weights)valid_ds = ConcatDataset(valid_ds, weights=dataset_weights)test_ds = ConcatDataset(test_ds, weights=dataset_weights)if random_mapping:world_size = torch.distributed.get_world_size(group=mpu.get_data_parallel_group())scale = max(200, 1 + (args.train_iters * args.batch_size * world_size) // len(train_ds))train_ds = RandomMappingDataset(train_ds, scale=scale)scale = max(200, 1 + ((1 + args.train_iters // args.eval_interval) * args.eval_iters * args.eval_batch_size * args.gradient_accumulation_steps * world_size) // len(valid_ds))valid_ds = RandomMappingDataset(valid_ds, scale=scale)test_ds = RandomMappingDataset(test_ds)return train_ds, valid_ds, test_ds

五、sat库之make_data_loader函数源码解读

在三说过类似torch的dataloader方法，在这里我大致说下，这里loader主要对dataset包装，使用顺序采用，配置类似的world_size与相应环境等内容，这也是库本身包装好的，可直接使用，其源码如下：

def make_data_loader(dataset, batch_size, args, split, collate_fn=None):world_size = torch.distributed.get_world_size(group=mpu.get_data_parallel_group())rank = torch.distributed.get_rank(group=mpu.get_data_parallel_group())distributed = world_size > 1# if IterableDataset, assume everything is properly configured. (pre-sharded) if isinstance(dataset, IterableDataset):if split in ['val', 'test'] and args.strict_eval:raise ValueError('IterableDataset cannot be used for validation or testing if `args.strict_eval=True`, because we cannot infer the length of the final batch before reading out them.')args.val_last_shape = [1] * world_size # just fake it, not actually usedargs.val_drop_number = 0args.test_last_shape = [1] * world_sizeargs.test_drop_number = 0return torch.utils.data.DataLoader(dataset,batch_size=batch_size//world_size,num_workers=args.num_workers,pin_memory=True,collate_fn=collate_fn,prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None,)sampler = torch.utils.data.SequentialSampler(dataset)  # 顺序采样drop_last = False # COMMENT: this is already solved by the complex logic of last_shape and drop_number.# the GPUs in the same model parallel group receive the same dataif distributed: # TODO reformat this, but it is not urgentgradient_accumulation_steps = getattr(args, 'gradient_accumulation_steps', 1)batch_sampler = DistributedBatchSampler(sampler,batch_size,drop_last,rank,world_size,gradient_accumulation_steps=gradient_accumulation_steps)else:batch_sampler = torch.utils.data.BatchSampler(sampler,batch_size,drop_last)last_len = len(dataset) % batch_sizebatch_per_worker = batch_size // world_sizelast_shape = [batch_per_worker] * (last_len//batch_per_worker) # some processes get full batchif last_len != 0:if last_len % batch_per_worker != 0:last_shape.append(last_len % batch_per_worker) # one process get the rest (<1 batch)drop_number = world_size - ((last_len-1)//batch_per_worker + 1)# other processes get nothing, but append 1 for running. will drop later according to drop_number.for j in range(drop_number): last_shape.append(1)else:drop_number = 0if split=='val':args.val_last_shape = last_shapeargs.val_drop_number = drop_numberelif split=='test':args.test_last_shape = last_shapeargs.test_drop_number = drop_numberdata_loader = torch.utils.data.DataLoader(dataset,batch_sampler=batch_sampler,num_workers=args.num_workers,pin_memory=True,collate_fn=collate_fn,prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None,)return data_loader

六、ItemDataset数据类源码解读

在这里，终于进入最核心部分，数据加工

class ItemDataset(Dataset):def __init__(self, image_processor, text_processor, args, data_dirs, cross_image_processor=None, **kwargs):super().__init__()self.data = self.load_data(data_dirs)  # 获得.jpg图片绝对路径的列表，保存到self.data变量中self.image_processor, self.text_processor, self.cross_image_processor = image_processor, text_processor, cross_image_processor # 传递的数据加工函数def process_img(self, img):img_dict = {'vision': self.image_processor(img)}if self.cross_image_processor:img_dict.update({'cross': self.cross_image_processor(img)})return img_dictdef process_text(self, answer, prompt):return self.text_processor(answer, prompt)def load_data(self, data_dir):all_files = find_all_files(data_dir, suffix=".jpg")print_rank0(f"find {len(all_files)} samples in all...")return all_filesdef __len__(self):return len(self.data)def __getitem__(self, index):data = self.data[index]  # 获得图片的绝对路径# imgtry:img = Image.open(data).convert('RGB')  # 载入图片except Exception as e:print_rank0(e, level=logging.WARNING)return {}img_dict = self.process_img(img)  # 图像加工# textlabel = data.split('/')[-1].split('.')[0]uni_key = labeltext_dict = self.process_text(label, "CAPTCHA:")if text_dict is None:print_rank0(f"Process text failed. Please check the max_target_length & max_source_length.\n The data is {data}", level=logging.WARNING)return {}# other attrret = {**img_dict, **text_dict, "question_id": uni_key}return ret

1、图像数据处理

1、process_img(self, img)

图像处理最终为一个img_dict的字典，包含2部分处理，其源码如下：

    def process_img(self, img):img_dict = {'vision': self.image_processor(img)}if self.cross_image_processor:img_dict.update({'cross': self.cross_image_processor(img)})return img_dict

2、self.image_processor(img)

image_processor函数传递为blip2_image_processor_func_with_inputs函数

def blip2_image_processor_func_with_inputs(image_processor, image):return {'image': image_processor(image).unsqueeze(0), 'input_ids': torch.zeros(1, 1, dtype=torch.long), 'position_ids': None, 'attention_mask': torch.ones(1, 1, dtype=torch.long)}

2、文本数据处理

1、原始图像验证码文本数据处理

原始验证码数据是一个图片，其名字为验证码数字命名，在源码data表示图像绝对路径(/home/*/0a4Ovs8789.jpg)。因此，使用以下方式label = data.split(‘/’)[-1].split(‘.’)[0]即可获得文本（0a4Ovs8789），随后使用self.process_text函数即可实现验证码文本数据。

label = data.split('/')[-1].split('.')[0]  
uni_key = label
text_dict = self.process_text(label, "CAPTCHA:")

验证码数据图如下：

2、使用自己文本数据代码修改

假如一个图片对应一个数据json文件，其内容如下：

{"conversations": [{"role": "assistant","content": "虽然无法从照片中确定他们的确切目的地或目的地，但这两名滑板运动员很可能正在使用公共交通工具前往滑板公园、休闲场所或其他可以练习滑板的地方。他们也可以在空闲时间携带滑板进行娱乐活动，往返于学校、工作或其他日常活动。或者，他们可能只是带着滑板在不同地点之间旅行，作为首选的交通方式。"}]
}

代码内容是获取json文件相应内容：

        json_path = data.replace('images','labels_zh')[:-4]+'.json'label=self.read_json(json_path)label = label['captions'][0]['content']  # 获取描述内容uni_key = labeltext_dict = self.process_text(label, "CAPTCHA:")

读取json辅助代码如下：

    def read_json(self,json_root):import jsonwith open(json_root, encoding='utf-8') as f:json_info = json.load(f)return json_info

最后文本将使用以下方式加工文本，我后期有时间在具体解读。

text_processor = llama2_text_processor(tokenizer, args.max_length, args.image_length)  # 获得文本加工函数