【SageMaker】Training JobsでFine-Tuningを行う際にmodel.tar.gzをS3から読み込む

以前の記事で、SageMaker Training JobsによるTabBERTモデルの事前学習を行ったので、今回は事前学習の結果`model.tar.gz`を元にFine-Tuningを実行するJobを作成しました。

基本的には事前学習と同じようなJobなのですが、以下の部分で工夫が必要だったのでメモとしてまとめました。

- tarファイルの展開
- 環境変数によるローカルとSageMaker間での引数の切替

https://qiita.com/suzuki0430/items/47228a8af26ccb4a88a2

なお、Fine-Tuningの元コードについては、以下の記事で解説しています。

https://qiita.com/suzuki0430/items/831dd5cad80c71bfe956

# tarファイルの展開
事前学習のJobによって、S3上に`model.tar.gz`が保存されています。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/a86bce14-a514-6cbc-6455-6b9b13acbae0.png)

`model.tar.gz`内には、モデル`pytorch_model.bin`、設定ファイル`config.json`、辞書ファイル`vocab.nb`、トークン→id変換ファイル`vocab_token2id.bin`が入っています。
Fine-Tuningではこれらを読み込む必要があるため、Jobを実行するときにtarファイルを展開するような工夫を行います。

まずは、Jobファイルの`input_model`で`model.tar.gz`のS3パスを指定します。
これでJob実行時に`model.tar.gz`が`/opt/ml/input/data/input_model/`(model_path)以下に置かれます。

```py:fine-tuning_jobs.ipynb
import sagemaker
from sagemaker.estimator import Estimator

session = sagemaker.Session()
role = sagemaker.get_execution_role()

estimator = Estimator(
    image_uri=<イメージURL>,
    role=role,
    instance_type="ml.g4dn.2xlarge",
    instance_count=1,
    base_job_name="tabformer-opt-fine-tuning",
    output_path="s3://<バケット名>/sagemaker/output_data/fine_tuning",
    code_location="s3://<バケット名>/sagemaker/output_data/fine_tuning",
    sagemaker_session=session,
    entry_point="fine-tuning.sh",
    dependencies=["tabformer-opt"],
    hyperparameters={
        "data_root": "/opt/ml/input/data/input_data/",
        "data_fname": "summary",
        "output_dir": "/opt/ml/model/",
        "model_path": "/opt/ml/input/data/input_model/",
    }
)
estimator.fit({
    "input_data": "s3://<バケット名>/sagemaker/input_data/summary.csv",
    "input_model": "s3://<バケット名>/sagemaker/output_data/pre_training/tabformer-opt-2022-12-16-07-00-45-931/output/model.tar.gz"
})
```

次にFine-Tuningの実行ファイル`tabformer_bert_fine_tuning.py`上に以下を記載します。

```py:tabformer_bert_fine_tuning.py
        with tarfile.open(name=path.join(args.model_path, f'model.tar.gz'), mode="r:gz") as mytar:
            mytar.extractall(path.join(args.model_path, f'model'))
            
            token2id_file = path.join(args.model_path, f"model/vocab_token2id.bin")
            vocab_file = path.join(args.model_path, f"model/vocab.nb")
            pretrained_model = path.join(args.model_path, f"model/checkpoint-500/pytorch_model.bin")
            pretrained_config = path.join(args.model_path, f"model/checkpoint-500/config.json")
```

`tarfile.open()`で`model.tar.gz`が読み込まれ、`mytar.extractall(path.join(args.model_path, f'model'))`で`/opt/ml/input/data/input_model/model/`以下に中身が展開されます。

これで、`token2id_file = path.join(args.model_path, f"model/vocab_token2id.bin")`のように展開されたファイルを読み込むことができるようになります。

# 環境変数によるローカルとSageMaker間での引数の切替
これでS3上の`model.tar.gz`を読み込めるようになったのですが、ローカルでFineTuningを行う際には読み込み先を変えたいケースもあると思います。

そんなときは、`os.getenv('SM_MODEL_DIR')`でSageMakerの環境変数`SM_MODEL_DIR`(コンテナ終了時にS3へアップロードされるディレクトリ)を取得し、ローカルとSageMaker(のJob)で読み込み先を切り替えます。

```py:tabformer_bert_fine_tuning.py
    key = os.getenv('SM_MODEL_DIR')
    
    if key :
        with tarfile.open(name=path.join(args.model_path, f'model.tar.gz'), mode="r:gz") as mytar:
            mytar.extractall(path.join(args.model_path, f'model'))
            
            token2id_file = path.join(args.model_path, f"model/vocab_token2id.bin")
            vocab_file = path.join(args.model_path, f"model/vocab.nb")
            pretrained_model = path.join(args.model_path, f"model/checkpoint-500/pytorch_model.bin")
            pretrained_config = path.join(args.model_path, f"model/checkpoint-500/config.json")
    else :
            vocab_file = path.join(args.model_path, f"vocab.nb")
            token2id_file = path.join(args.model_path, f"vocab_token2id.bin")
            pretrained_model = path.join(args.model_path, f"checkpoint-500/pytorch_model.bin")
            pretrained_config = path.join(args.model_path, f"checkpoint-500/config.json")
```

# 参考資料

https://koshian2.booth.pm/items/4150087

https://hibiki-press.tech/python/tarfile/1435

https://www.lifewithpython.com/2021/09/python-archive-files.html