4
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

記事投稿キャンペーン 「2024年!初アウトプットをしよう」

Huggingfaceのdiffusersでpretrainedのモデルの保存先を変更する

Posted at

Huggingfaceのdiffusersではpipeline.from_pretrainedを実行することでHuggingface Hubからpretainedのモデルをダウンロードできます

ダウンロードしたモデルは~/.cache/huggingfaceにキャッシュとして保存されます

Docker環境ではキャッシュ先のディレクトリをマウントしておかないとコンテナを破棄したタイミングでキャッシュが失われてしまいます

今回はキャッシュ先のディレクトリを~/.cache/huggingfaceから<work_dir>/.cache/huggingfaceに変更してみます

試してみたところ,キャッシュ先のディレクトリは環境変数$HF_HUB_CACHEで指定されています.$HF_HUB_CACHE$HF_HOME/hubに設定されているので,環境変数$HF_HOMEを変えれば良さそうです

    export HF_HOME=<work_dir>/.cache/huggingface

Docker環境ではENVコマンドで環境変数を設定できます

    FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04
    LABEL maintainer="Hugging Face"
    LABEL repository="diffusers"
    
    # ワークディレクトリの指定
    WORKDIR /app

    # 非インタラクティブモードにする (入力待ちでブロックしなくなる)
    ENV DEBIAN_FRONTEND noninteractive
    # .pycを作らないように
    ENV PYTHONDONTWRITEBYTECODE 1
    # バッファの無効化
    ENV PYTHONUNBUFFERED 1
    # torchvisionでpretrainedのモデルを保存する場所
    ENV TORCH_HOME /app/.cache/torchvision
    
    # ----------------------------------------------------------------
    # setup (root) 
    # ----------------------------------------------------------------
    RUN apt update && \
        apt install -y bash \
        build-essential \
        git \
        git-lfs \
        curl \
        ca-certificates \
        libsndfile1-dev \
        libgl1 \
        python3.8 \
        python3-pip \
        python3.8-venv && \
        rm -rf /var/lib/apt/lists
    
    # ----------------------------------------------------------------
    # create user
    # ----------------------------------------------------------------
    # UIDとGIDは外から与える
    ARG USER_UID
    ARG USER_GID
    
    # コンテナ内でのユーザー名, グループ名
    ARG USER_NAME=user
    ARG GROUP_NAME=user
    
    # グループが存在しなかったら, 適当なグループを作成
    RUN if ! getent group $USER_GID >/dev/null; then \
        groupadd -g $USER_GID $GROUP_NAME; \
        fi
    
    # ユーザーを作成
    RUN useradd -m -u $USER_UID -g $USER_GID -s /bin/bash $USER_NAME
    
    # 初期ユーザーの変更
    USER $USER_NAME
    
    # ----------------------------------------------------------------
    # setup (user) 
    # ----------------------------------------------------------------
    
    # pre-install the heavy dependencies (these can later be overridden by the deps from setup.py)
    RUN python3 -m pip install --no-cache-dir --upgrade pip && \
        python3 -m pip install --no-cache-dir \
        torch \
        torchvision \
        torchaudio \
        invisible_watermark && \
        python3 -m pip install --no-cache-dir \
        accelerate \
        datasets \
        hf-doc-builder \
        huggingface-hub \
        Jinja2 \
        librosa \
        numpy \
        scipy \
        tensorboard \
        transformers \
        omegaconf \
        pytorch-lightning \
        xformers
    
    ENV HF_HOME /app/.cache/huggingface
    
    CMD ["/bin/bash"]

終わり

4
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
4
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?