adm显卡下使用gpu尝试

一、定义

1. windows安装使用教程
2. 注意事项，兼容性
3. 解决torch 版本低，transformer加载模型报错

二、实现

1. windows 安装使用教程

   1.1 检查显卡

   win+r输入dxdiag，到显示里面找设备名称，我的是AMD显卡，不适合使用torch-gpu版本。

1.2 安装配置

1. pip install torch-directml        
#注意，需要在python 10以及其版本以下，不然不支持
# torch-directml 下安装的torch 目前最高版本为2.4， transfomer可能不支持
# torch-directml 相当于一个插件，有些torch 算子是不支持的。如torch.kl_div()
2. 测试import torchimport torch_directml# 获取可用的 DirectML 设备（你的AMD显卡）dml_device = torch_directml.device()# 将你的Tensor和模型放到这个设备上tensor = torch.rand(2, 2).to(dml_device)

2. 注意事项

   adm 环境下训练的模型不能直接使用在nvidia下，需要进行模型迁移。

   方法一、转为onnx格式； 方法二：保存state_dict。torch.save(model.state_dict(), ‘model_weights.pth’)

3. 解决torch 版本低，transformer加载模型报错

   Due to a serious vulnerability issue in torch.load, even with weights_only=True, we now require users to upgrade torch to at least v2.6 in order to use the function. This version restriction does not apply when loading files with safetensors.

解决：使用safetensor 方式加载。

teacher_model = AutoModelForSequenceClassification.from_pretrained(config.teacher_model_name,use_safetensors = True
)