# SplitComputing **Repository Path**: HouEna/split-computing ## Basic Information - **Project Name**: SplitComputing - **Description**: kkkksdadasdas - **Primary Language**: Unknown - **License**: Not specified - **Default Branch**: master - **Homepage**: None - **GVP Project**: No ## Statistics - **Stars**: 0 - **Forks**: 0 - **Created**: 2025-10-10 - **Last Updated**: 2026-01-06 ## Categories & Tags **Categories**: Uncategorized **Tags**: None ## README # Hn-Split ## 模型配置 在配置文件中，tercher_model或者student_model 下的name、params 是获取模型的关键参数对 这个name 可以是一个函数名字例如“get_timm_model” 然后搭配着params来获取对应的模型 torchdistill.models.registry 是在此框架下得到模型的关键库 其中有两个自定义模型方式： 1. 用修饰符 @register_model_class 直接注册模型 2. 用修饰符 @register_model_func 注册获取模型的函数 在调用get_model的时候会从register_model_class、register_model_func中找是否存在model_name，存在的话就把paramer传递给这个模型、或者构建模型的函数，从而得到一个模型 如果不存在那么就从torch.hub中通过传入的网址加载 例如: ```yaml teacher_model: name: &teacher_model_name 'get_timm_model' params: timm_model_name: &backbone_base 'swin_s3_base_224' no_classes: 1000 reset_head: False pretrained: True assign_layer_names: False split_idx: -1 ``` get_timm_model就是用@register_model_func注册的函数，params里面的参数就是这个函数接收的参数 ## 评估指标(eval metric)配置 在配置文件config_yaml中train的配置中要有训练的阶段stage1、stage2、。。。 每个训练阶段都要有一个eval_metrics属性，用来配置在这个阶段训练的时候用什么评估指标 每个阶段的评估指标存在stage_eval_metrics这个字典中， 所有阶段的评估指标字典放在eval_metrics这个列表下 例如 ```yaml train: log_freq: 500 epoch_to_update: &epoch_to_update 5 stage1: eval_metrics: [ 'accuracy', 'bpp' ] num_epochs: *epoch_to_update train_data_loader: dataset_id: *imagenet_train random_sample: True batch_size: 16 num_workers: 4 cache_output: val_data_loader: dataset_id: *imagenet_val random_sample: False batch_size: 128 num_workers: 4 ``` 这些指标会通过`get_eval_metric`来获取对应的计算对象,可支持的计算方法定义在EVAL_METRIC_DICT中 ```python EVAL_METRIC_DICT = { "accuracy": EvaluationMetric(eval_func=evaluate_accuracy, init_best_val=0, comparator=lambda curr_top_val, epoch_val: epoch_val > curr_top_val), "bpp": EvaluationMetric(eval_func=evaluate_bpp, init_best_val=float("inf"), comparator=lambda curr_top_val, epoch_val: epoch_val < curr_top_val), 'psnr': EvaluationMetric(eval_func=evaluate_psnr, init_best_val=float("-inf"), comparator=lambda curr_top_val, epoch_val: epoch_val > curr_top_val), "accuracy-and-filesize": EvaluationMetric(eval_func=evaluate_filesize_and_accuracy, init_best_val=None, comparator=None ) } ``` 是一个类对象EvaluationMetric ### 多阶段配置 #### models部分 如果是多阶段，那么第二阶段的时候多阶段蒸馏框架回去获取config.yaml里面的models部分中的teacher_model_{阶段号} 代码中是这样的： ```python #到了第二阶段 if _is_multi_stage(training_box) and training_box.stage_number == 2: logger.info(f"Finished Stage {curr_stage}..") #获取第二个阶段的模型配置 new_teacher_config = config.get("models").get(f"teacher_model_{training_box.stage_number + 1}") ``` 如果是多阶段，那么就会将第一个阶段的模型权重存储在config.yaml的student_model中`ckpt_stage1`所指定的路径下 ## DistillationBox的逻辑 ### 初始化 #### 初始化入口 实例化相关对象的时候从 __init__进入，然后调用self.setup()进行进一步的初始化 #### 初始化dataloader 调用setup_setup_data_loaders，从train_config拿到train_data_loader 和 val_data_loader的配置，初始化self.train_data_loader 和 self.val_data_loader #### 初始化模型 1. 调用setup_teacher_student_models 从train_config拿到teacher 和 student 的配置，这是教师模型、学生模型的配置（必须是teacher 和 student 这两个字段，而不是teacher_model 等，因为库识别的就是这俩字段） 调用redesign_model根据教师模型、学生模型中的frozen_modules字段，冻结网络对应的部分，根据sequential字段重新构建在蒸馏中两个模型需要的部分。 冻结、提取完需要的部分后的网络存储在 self.teacher 和 self.student中 2. 调用set_hooks函数根据train_config中teacher、student中的forward_hook字段构建pairs，pairs的键是layer的名，值是layer的函数 最终所有的pair都保存在self.target_teacher_pairs和self.target_student_pairs中。也就是说这里实际上保存的是我们希望获取到的层输入、输出和这个输入输出对应的获取函数 3. 根据train_config中teacher和student字段中的requires_grad字段，如果为true，则冻结的只是frozen_modules中选择的模型部分(已经在setup_teacher_student_models中处理过了) 如果为False，则直接冻结模型所有权重 #### 初始化损失函数 调用setup_loss函数 , 根据criterion_config 将loss类存在 #### data_loader 调用的是setup_data_loaders() 1. 根据yaml配置文件中的`train_data_loader`和`val_data_loader`字段得到train、val dataloader的config 2. 调用函数build_data_loaders ### 内置方法 #### get_teacher_output 1. 调用self.teacher_forward_proc,传入教师模型、batch数据、target，得到教师模型的输出teacher_outputs 其中self.teacher_forward_proc（self.student_forward_proc）的工作原理如下： * 它是在初始化阶段在setup_teacher_student_models的时候调用函数get_forward_proc_func，传入教师、学生模型的config中的forward_proc字段 如果没有forward_proc字段，那么就传入None，如果传入None，那么返回的self.teacher_forward_proc就是函数forward_batch_only： ```python @register_forward_proc_func def forward_batch_only(model, sample_batch, targets=None, supp_dict=None): return model(sample_batch) ``` 与之相对应的还有集中前向转播的函数，也就是说如果在config里面forward_proc字段中有`forward_batch_target`、`forward_batch_supp_dict` `forward_batch4sskd`值，就会在前向传播的时候以对应的方式传入数据。实际上forward_proc字段是定义模型的前向传播的传参格式 ```python @register_forward_proc_func def forward_batch_target(model, sample_batch, targets, supp_dict=None): return model(sample_batch, targets) @register_forward_proc_func def forward_batch_supp_dict(model, sample_batch, targets, supp_dict=None): return model(sample_batch, supp_dict) @register_forward_proc_func def forward_batch4sskd(model, sample_batch, targets=None, supp_dict=None): c, h, w = sample_batch.size()[-3:] sample_batch = sample_batch.view(-1, c, h, w) return model(sample_batch) ```` 2. 调用 extract_io_dict函数，传入self.teacher_io_dict,得到提取到的extracted_teacher_io_dict 3. 返回 teacher_outputs（输入的数据经过教师模型的所有层得到的最终的输出），extracted_teacher_io_dict（感兴趣的中间层的输出） ### forward 1. 调用get_teacher_output,传入的是batch数据以及对应的target，得到teacher_outputs（输入的数据经过教师模型的所有层得到的最终的输出）和extracted_teacher_io_dict（感兴趣的中间层的输出） 2. 调用self.student_forward_proc得到student_outputs（输入的数据经过学生模型的所有层得到的最终的输出），调用extract_io_dict得到extracted_student_io_dict（感兴趣的中间层的输出） 3. extracted_teacher_io_dict和extracted_student_io_dict组成字典output_dict ~~~ output_dict = {'teacher': extracted_teacher_io_dict, 'student': extracted_student_io_dict} ~~~ 4. 将output_dict传入self.criterion得到total_loss ## CustomLoss类 1. 根据criterion_config中的sub_terms，将每一项的criterion作为sub_criterion_config,然后调用函数get_single_loss获取到具体的sub_criterion,存入term_dict中 即self.term_dict存储的是yaml配置文件中criterion字段中sub_terms字段的内容解析之后得到的对应loss对象 ```yaml criterion: type: 'GeneralizedCustomLoss' org_term: factor: 0.0 sub_terms: layer1: criterion: type: 'MSELoss' params: reduction: 'sum' params: input: is_from_teacher: False module_path: 'compression_module' io: 'output' target: is_from_teacher: True module_path: 'layers.1' io: 'output' factor: 1.0 bpp: criterion: type: 'BppLossOrig' params: input_sizes: *input_size entropy_module_path: 'compression_module.entropy_bottleneck' reduction: 'sum' factor: *lmbda ``` 最终term_dict中保存的是{‘layer1’:对应的Loss对象，‘bpp’:对应的loss对象} 这些loss对象是根据criterion中的type去对应的注册的loss列表中找到的，而这些loss对象的初始化 是根据params参数传入对应的对象完成的。 上面的这个例子中可以看到layer1这个item有两个params，我们通过get_single_loss这个函数去找的时候是找第一层的params作为参数，如果 第一层没有params，那么loss_config.get('params', None)就会为None，那么get_single_loss就知道了要到SINGLE_LOSS_CLASS_DICT 中找这个loss对象，这个bpp就是这样的 而像MSELoss这个loss本身是一个torch.nn.modules.loss自带的一个loss对象，里而不是第三方注册的 注意：criterion下面的params是真正的作用与type中loss对象的配置参数，而与criterion同一层级的params是用来对loss对象进一步进行包装 使用的配置参数。 MSELoss是torch自带的loss 但是torchdistll为了更方便地用于知识蒸馏，因此用get_loss_wrapper来传递与criterion同一层级的params来 进行更适合于蒸馏的时候使用地loss对象 这个最终用到了torchdistll的SimpleLossWrapper对象 ## MultiStagesDistillationBox使用 ### forward() 1. 输入: ### post_process()方法 调用post_process方法，进入下一个阶段 # 参数 * --skip_ckpt: 如果带了 就会加载模型的时候跳过加载权重，没带的话就会从config.yaml里面的model下面的ckpt项规定的权重路径 * --reset_student_hand:如果分类的数目变了，就用它来设置分多少类 * --validation_metric:用来指导保存权重的指标，它规定了在对比模型的好坏的时候由哪个指标决定好与坏，更好的权重将被保存 * --moco_train：使用moco的方式训练学生模型编码器，要搭配MocoinfoNCELoss，以及MoCoEncoder使用。其中负样本队列的容量由参数--moco_queue_size 来控制，默认是4096 * --moco_proj_dim：moco训练时候的哪个映射头的维度，默认是128， 注意这个如果要改动的话，config文件中的encoder的proj_dim要和它一样才行 # EntropyBottleneck的效果 EntropyBottleneck的作用有两个： 1. 量化：把连续值y变成离散整数$\hat{y}$ 2. 估计概率：估计出每个整数值$\hat{y}$的概率$p(\hat{y})$,用于计算比特数$-log_2p(\hat{y})$ ## 如何计算$p(\hat{y})$ ### 计算 logits * logits是通过一个可学习网络学出来的 * EntropyBottleneck的self.filters规定了这个可学习网络的隐藏层的输出维度 * 经过转换，最终这个可学习网络的输入，输出维度是输入为1，输出为1，中间的为filters 的规定，默认是3，3，3，3。 filter的通道数和输入的数据的通道数相同。 * 这些用来学习的权重、偏置矩阵通过register_parameter注册到对应的参数属性中： ~~~ self.weight = nn.Parameter(torch.randn(10, 10)) 等价于： w = nn.Parameter(torch.randn(10, 10)) self.register_parameter('weight', w) 它们 本质一样：把参数加入模块的 ._parameters 字典里。 EntropyBottleneck中： self.register_parameter(f"_matrix{i:d}", nn.Parameter(matrix)) self.register_parameter(f"_bias{i:d}", nn.Parameter(bias)) self.register_parameter(f"_factor{i:d}", nn.Parameter(factor)) 也就是把可学习的权重矩阵存在了 self._matrix0、1、2、3... self._bias0、1、2、3... self._factor0、1、2、3..中 ~~~ * _logits_cumulative()使用self._matrix与输入矩阵相乘再加上self._bias,再加上self._factor做相应的变化的值得到最终的logit： ~~~ logits = inputs logits = torch.matmul(F.softplus(matrix), logits) softplus(z)=log(1+e^z)保证输出永远非负，为的是保持单调性，只要权重非负就不会破坏单调性。 matmul就是矩阵乘法会把最前面的维度 C 当成 batch‑dim，在每个通道里做标准矩阵乘。 logits += bias logits += torch.tanh(factor) * torch.tanh(logits) ~~~ ### 由logits得到CDF,CDF做差得到概率p CDF是累计分布函数，是概率密度函数PDF的积分，CDF(x)=P(X<=x),也就是随机变量X小于等于x的概率是多少 logits(p) = log(p/(1-p))这个变化是为了让网络输出在正负无穷区间上做回归，再通过sigmoid拉回[0,1] _logits_cumulative()是CDF的logit，然后通过sigmoid得到CDF PDF是CDF的导数，即：$p(y)=\frac{dCDF(y)}{dy}$ 也就是$CDF(y_1) - CDF(y_0) = P(y_0 < Y < y_1)$ 我们要求的量化后的 $p(\hat{y})$实际上就是p(y-0.5