Ｗindows 下 AMD显卡训练模型有救了：pytorch_directml 下运行Transformers

原创

已于 2023-02-21 15:54:16 修改 · 4.6k 阅读

11 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #人工智能

于 2023-02-21 09:00:02 首次发布

本文介绍了在Windows操作系统下，使用AMD显卡通过PyTorchDirectML插件训练Transformer模型的方法，特别是针对codebert的设置。由于pipeline可能存在的兼容性问题，建议直接编写代码来执行模型训练，以确保GPU得到有效利用。

Windows 下ａｍｄ显卡训练transformer 模型。安装方法参见： Windows下用amd显卡训练： Pytorch-directml 重大升级，改为pytorch插件形式，兼容更好_amd显卡 pytorch_znsoft的博客-优快云博客

import os
import imp
try:
    imp.find_module('torch_directml')
    found_directml = True
    import torch_directml
except ImportError:
    found_directml = False

import torch
from transformers import RobertaTokenizer, RobertaConfig, RobertaModel, RobertaForMaskedLM,pipeline

DIR="E:/transformers"
MODEL_NAME="microsoft/codebert-base"
from transformers import AutoTokenizer, AutoModel

if found_directml:
    device=torch_directml.device()
else:
    device=torch.device("cpu")

# tokenizer = AutoTokenizer.from_pretrained(DIR+os.sep+MODEL_NAME)
# model = AutoModel.from_pretrained(DIR+os.sep+MODEL_NAME).to(device)
# nl_tokens=tokenizer.tokenize("return maximum value")