# fireoil

**Repository Path**: zhangxinpaper/fireoil

## Basic Information

- **Project Name**: fireoil
- **Description**: 我的代码模板，可以应用于各种训练场景，只需要实现dataset, model, compute_loss和compute_metric即可，支持多机多卡
- **Primary Language**: Python
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 4
- **Forks**: 1
- **Created**: 2022-07-07
- **Last Updated**: 2024-06-07

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 使用自定义pytorch训练框架

​	做这个仓库的原因是为了解决Pytorch训练代码重复编写和代码灵活度之间矛盾的问题。在日常工作中，经常遇到要重复写训练代码，且代码量又很大的问题，一般算法工程师会有一套自己的模板代码。但是这个代码维护麻烦，且通用性一般不足，然后如果使用fastai, detectron2, transformers, pytorch_lightning了，容易出现不知名的BUG，官方代码查起来由很复杂，因此萌生了编写自定义训练代码的冲动，于是有了下面的代码仓库，目前开源在gitee上，github因为众所周知的原因没有上传。

​	目前代码支持如下功能。（具体参数名称，可以查看Trainer的文档或代码注释部分）

- EMA模型训练
- 指定网络名称字符进行对抗训练
- DistributedDataParallel分布式训练
- 梯度裁剪
- 支持CNN和Transformer模型参数分层设置weight decay
- 混合训练
- 最大范数约束
- 梯度累积
- warmup学习率

> gitee网址：<https://gitee.com/zhangxinpaper/fireoil>

## **使用方法**

- 将fireoil文件夹拷贝到目录下

- 编写代码

  - Dataset，定义数据集

    - train_ds: 训练数据dataset
    - val_ds: 验证数据dataset
    - 然后完全自主的pytorch数据定义方式编写Dataset代码吧！！！

  - Model

    - 完全自主的pytorch模型定义方式

  - Trainer，继承Trainer

    - 编写compute_loss和compute_metrics函数

    ### **编写compute_loss函数**

    - ```python
      def compute_loss(self, model, batch, device, loss_fn=None, return_preds=False):
              return {"loss": loss, "preds": preds} if return_preds else {"loss":loss}
      ```

    - 输入输出的格式不要变化，给loss赋值，进行梯度backward和模型权重更新，preds保存，其会原封不动的作为compute_metrics的参数

    - model是前面定义的模型，Trainer自动将model传输到device中。batch是dataset经过DataLoader和data_collator处理后批数据，这里需要将batch也放到相应的device中，loss_fn是初始化Trainer导入的，也可以在model内部自定义，那么loss_fn就可以不使用了。return_preds保留，不要修改它。

    ### **编写compute_metrics**

    - ```python
      def compute_metrics(self, preds, batch):
              return {"metric": metric, "other": other, ...} # metric和other必须是平均之后的
      ```

    - preds是compute_loss函数的输出，batch是批量处理过的数据，因为指标默认使用cpu操作，因此，需要将preds转换到cpu上操作;

    - metric和other都是函数运行的结果，metric是作为模型保存的指标;

    - compute_metrics的返回结果将会以pd.DataFrame和tensorboardX的格式保存，在self.args.saved_model_path中可以看到，其中saved_model_path是Trainer的传入args中设置，如果没有设置，则会使用默认路径，即与运行代码在同一层级的trained_weights文件夹中。



### 一个例子

- Args

```bash
# 默认参数设置
'debug': True,
'seed': 402,
'deterministic': True,
'gpus': 1,
'train_batch_size_per_device': 16,
'val_batch_size_per_device': 32,
'num_workers': 4,
'use_model_ema': True,
'sync_bn': False,
'resume': None,
'use_transformer_arch': True,
'learning_rate': 2e-05,
'weight_decay': 0.01,
'momentum': 0.9,
'warmup_ratio': 0.2,
'gradient_accumulate_steps': 1,
'num_epochs': 6,
'use_clip_norm': True,
'max_norm': 10.0,
'use_mixed_precision_training': False,
'use_fgm_training': False,
'fgm_emb_str': 'emb.',
'logdir': None,
'checkpoint_style': 'max',
'saved_model_path': './trained_weights'
```

```python
# 参数设置，根据上面的参数key进行调整
def make_parser():
    parser = ArgumentParser(description="Image classifier")
    parser.add_argument("--use_transform_arch", default=False, type=bool)
    parser.add_argument("--num_workers", default=0, type=int)
    parser.add_argument("--learning_rate", default=1e-3, type=float)
    parser.add_argument("--use_clip_norm", default=False, type=bool)
    parser.add_argument("--max_norm", default=1., type=float)
    parser.add_argument("--use_fgm_training", default=False, type=bool)
    parser.add_argument("--fgm_emb_str", default="conv1.", type=str)
    parser.add_argument("--gradient_accumulate_steps", default=1, type=int)
    parser.add_argument("--num_epochs", default=6, type=int)
    return parser

args = make_parser().parse_args([])
```

- Dataset

```python
class MyDataset(Dataset):
    def __init__(self, df, transform, parent_path="./ReversoContextClass/images/"):
        super().__init__()
        self.df = df
        self.transform = transform
        self.parent_path = parent_path
    
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        image_path = row["name"]
        label = row["label_idx"]
        
        data = Image.open(os.path.join(self.parent_path, image_path)).convert("RGB")
        data = self.transform(data)
        
        return data, torch.tensor(label, dtype=torch.long)
```

- **Model**

```python
model = resnet18(pretrained=False)
```

- Trainer

```python
class MyTrainer(Trainer):
    def compute_loss(self, model, batch, device, loss_fn=None, return_preds=False):
        x, y = batch # 对应Dataset中的data, torch.tensor(label, dtype=torch.long)
        x, y = x.to(device), y.to(device)
        y_hat = model(x)
        loss = loss_fn(y_hat, y)
        preds = y_hat
        
        return {"loss": loss, "preds": preds} if return_preds else {"loss":loss}

    def compute_metrics(self, preds, batch):
        x, y = batch        
        y_hat = preds.detach().cpu().argmax(dim=-1).numpy()
        y = y.detach().cpu().numpy()
        
        p, r, fb, _ = precision_recall_fscore_support(y, y_hat, average="macro")
        accuracy = accuracy_score(y, y_hat)
        
        metric = accuracy
        
        return {"metric": metric, "precision":p, "recall": r, "fbeta": fb} # 平均值
    
trainer = MyTrainer(args, model, train_ds, val_ds, loss_fn)
trainer.train() # 进行训练
```