关于模型训练不能跑满显卡

[ykallan]

import json
import argparse
from fengshen import UbertPipelines
import torch

total_parser = argparse.ArgumentParser("TASK NAME")
total_parser = UbertPipelines.pipelines_args(total_parser)
args = total_parser.parse_args()

args.pretrained_model_path = r'D:\nlp_about\pretrained_model\Erlangshen-Ubert-110M-Chinese' 
args.default_root_dir = './train/'  
args.max_epochs = 5
args.gpus = 1
args.batch_size = 768

data_path = r"D:\nlp_about\Fengshenbang-LM\datas\__data.json"

with open(data_path, "r", encoding="utf-8") as f:
    all_data = json.load(f)

total_length = len(all_data)

train_data = all_data[:40000]
test_data = all_data[40000:45000]

model = UbertPipelines(args)

model.fit(train_data, test_data)

不管我的 args.batch_size设置的多大，我是用的3090训练，显存只有11g左右，cuda跑在80%左右，也没有跑满，不确定是什么原因

[suolyer]

我没有遇过这个问题，batch_size是pytorch_lightning.trainer的一个超参数，你可以试试用这个trainer训其他模型是否出现同样的问题。

另外，我们对ubert做了一个很大的升级，可以尝试使用我们的升级版https://github.com/IDEA-CCNL/Fengshenbang-LM/tree/main/fengshen/examples/uniex