pytorch训练模型内存溢出

1、训练模型命令命令

如下所示是训练命名实体识别的命令，在win10系统下执行

activate pytorch
cd F:\Python\github\ultralytics-main\submain\pytorch_bert_bilstm_crf_ner-main
f:
python main.py --bert_dir="../model_hub/chinese-bert-wwm-ext/" --data_dir="./data/cner/" --data_name="cner" --model_name="bert" --log_dir="./logs/" --output_dir="./checkpoints/" --num_tags=33 --seed=123 --gpu_ids="0" --max_seq_len=150 --lr=3e-5 --crf_lr=3e-2 --other_lr=3e-4 --train_batch_size=32 --train_epochs=3 --eval_batch_size=32 --lstm_hidden=128 --num_layers=1 --use_lstm="False" --use_idcnn="True" --use_crf="True" --dropout_prob=0.3 --dropout=0.3

2、内存溢出

本机GPU是6G的，所需明显超出本机显存，如下图

3、解决方法

尝试减少训练批次大小（train_batch_size）来减少GPU内存的使用。可以将train_batch_size的值从32减少到16或更小的值

activate pytorch
cd F:\Python\github\ultralytics-main\submain\pytorch_bert_bilstm_crf_ner-main
f:
python main.py --bert_dir="../model_hub/chinese-bert-wwm-ext/" --data_dir="./data/cner/" --data_name="cner" --model_name="bert" --log_dir="./logs/" --output_dir="./checkpoints/" --num_tags=33 --seed=123 --gpu_ids="0" --max_seq_len=150 --lr=3e-5 --crf_lr=3e-2 --other_lr=3e-4 --train_batch_size=16 --train_epochs=3 --eval_batch_size=32 --lstm_hidden=128 --num_layers=1 --use_lstm="False" --use_idcnn="True" --use_crf="True" --dropout_prob=0.3 --dropout=0.3

4、结果

可以开始训练了