如何精简SBERT的sentence-transformer库？CPU与GPU环境优化指南

SBERT的sentence-transformers库是Python开源LLM生态中最流行的文本片段向量嵌入生成库，它拥有简洁的API，但依赖包体积异常庞大。

以下是构建包含该工具的基础Docker镜像的情况：构建耗时超11分钟，最终镜像大小达7.5GB：

参考基础Dockerfile内容：

FROM python:3.11.2-slim-bullseye
RUN pip install --upgrade pip && pip install sentence-transformers

原本以为镜像已预打包模型，但运行官方入门代码片段：

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('all-MiniLM-L6-v2')
embeddings = model.encode("vectorize this text")

时，程序又额外下载了数百MB文件。推测体积过大主要源于底层torch依赖（占6.9GB），而torch关联的nvidia-*依赖安装位置如下：Debian/Ubuntu系统的/usr/lib/x86_64-linux-gnu/目录，或RHEL/CentOS系统的/usr/lib64/目录；Torch自带的CUDA库则位于Python环境的site-packages/torch/lib目录下。

CPU环境精简依赖（已本地下载all-MiniLM-L6-v2模型） #

若已将HuggingFace的all-MiniLM-L6-v2模型下载至本地路径path/to/all-MiniLM-L6-v2，仅需在CPU环境运行代码，可通过以下方式实现精简：

• 安装CPU专属版Torch：默认pip安装sentence-transformers会拉取带CUDA支持的Torch，指定CPU版本可大幅削减体积：

  pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
  

• 手动安装核心依赖：先安装sentence-transformers运行必需的基础依赖，再跳过依赖检查安装库本身：

  pip install --upgrade pip
  # 安装CPU版Torch
  pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
  # 安装核心依赖：transformers、sentencepiece等
  pip install transformers sentencepiece numpy scipy scikit-learn tqdm
  # 安装sentence-transformers（跳过自动依赖安装）
  pip install sentence-transformers --no-deps
  

• Docker镜像多层优化：使用多阶段构建并清理缓存，减少镜像冗余：

  FROM python:3.11.2-slim-bullseye
  WORKDIR /app
  RUN pip install --upgrade pip \
      && pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu \
      && pip install transformers sentencepiece numpy scipy scikit-learn tqdm \
      && pip install sentence-transformers --no-deps \
      && rm -rf /root/.cache/pip
  # 复制本地模型到镜像内
  COPY path/to/all-MiniLM-L6-v2 /app/all-MiniLM-L6-v2
  

• 加载本地模型避免下载：代码中直接指定本地模型路径，无需在线拉取：

  from sentence_transformers import SentenceTransformer
  model = SentenceTransformer('/app/all-MiniLM-L6-v2')
  embeddings = model.encode("vectorize this text")
  

GPU环境无冗余安装的边际依赖 #

切换至GPU环境运行时，在CPU精简方案的基础上，仅需添加以下依赖：

• 匹配CUDA版本的Torch：根据NVIDIA驱动版本选择对应CUDA版本的Torch，例如CUDA 11.8版本：

  pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
  

• NVIDIA运行时基础镜像（Docker环境）：无需手动安装复杂的NVIDIA驱动依赖，直接基于官方runtime镜像构建：

  FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
  RUN apt-get update && apt-get install -y python3 python3-pip \
      && rm -rf /var/lib/apt/lists/*
  WORKDIR /app
  RUN pip install --upgrade pip \
      && pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 \
      && pip install transformers sentencepiece numpy scipy scikit-learn tqdm \
      && pip install sentence-transformers --no-deps \
      && rm -rf /root/.cache/pip
  COPY path/to/all-MiniLM-L6-v2 /app/all-MiniLM-L6-v2
  

• GPU可用性验证：可添加代码确认Torch识别GPU：

  import torch
  print(torch.cuda.is_available())  # 输出True表示GPU可用
  

内容的提问来源于stack exchange，提问作者nmurthy