【AIGC】AIGC程序记录_automodel.from_pretrained 指定显卡

一、模型相关

二、程序相关

1、GPU使用

方式一：

os.environ["CUDA_VISIBLE_DEVICES"] = "3,4,5"   # 如果不生效，还是默认使用gpu0，将该句命令放到import torch前，或者放到文件最前面就可以了。

方式二：

1
2
3
4
5
6

2、是否加载GPU

 -- GPU加载：
 model = AutoModelForCausalLM.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat",
        cache_dir="./models"     # 默认下载的模型路径
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True,
    )
    
 -- 先CPU，再GPU加载：
 model = AutoModelForCausalLM.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat",
        cache_dir="./models"
        torch_dtype=torch.float16,
        trust_remote_code=True,
    )
 model = model.cuda()
 
 -- 量化:
 model = model.quantize(4).cuda()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

三、服务相关

1、chat类-websocket服务

from fastapi import FastAPI, Request, WebSocket
from transformers import AutoTokenizer, AutoModel
import uvicorn, json, datetime
import torch
import os
import asyncio

""" BaiChan13B Example """

#配置GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "3,4,5"
app = FastAPI(debug=False)

@app.websocket("/chat")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    messages = []
    while True:
        prompt = await websocket.receive_text()
        messages.append(
            {"role": "user", "content": prompt_text}
        )
        for response in model.chat(tokenizer, messages, stream=True):
            result = {
                "res":{
                    "text": response,
                    "end": False
                }
            }
            result = json.dumps(result, ensure_ascii=False)
            await websocket.send_text(result)
            await asyncio.sleep(0.001)
        result = {
            "res": {
                "text": response,
                "end": True
            }
        }
        result = json.dumps(result, ensure_ascii=False)
        await websocket.send_text(result)
        await asyncio.sleep(0.01)

        messages.append(
            {"role": "assistant", "content": response}
        )


if __name__ == "__main__":

    model = AutoModelForCausalLM.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat",
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True,
    )
    model.generation_config = GenerationConfig.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat"
    )
    tokenizer = AutoTokenizer.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat",
        use_fast=False,
        trust_remote_code=True
    )

    model.eval()
    uvicorn.run(app, host="0.0.0.0", port=8000, ws_ping_timeout=5)
        
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

注意

如果websocket返回结果比较慢，加入sleep，如下：

await websocket.send_text(result)
await asyncio.sleep(0.001)
1
2