🐥
Segformerを使って車のSemantic Segmentationを実施
はじめに
Semantic SegmentationのタスクにTransformerを使用した場合の性能を試してみたかったため、Segformerを使用してみました。
なお本リポジトリはKaggleのcarvana-image-masking-challengeを学習データとして学習させた際のソースコードになります。
今回作成したソースコードは以下のリポジトリに公開しているので、よかったら試してみてください。
参照
環境構築
SegformerはDocker環境に構築することにしました。理由は既存の環境を壊したくなかったため。また、もともとCUDA10.2が入っていたため、CUDA10.2上で環境を構築することにしました。
CUDA10.2のDockerのベースイメージをダウンロード
まず、ベースとなるnvidia/cuda:10.2-cudnn7-devel-ubuntu18.04を引っ張ってくる必要があるのですが、24年1月時点でEOLとなったCUDAバージョンのDockerイメージが削除されていたため、引っ張ってくる必要があります。CUDA10.2のDocker Imageを作成を参考にしてDocker Imageを作成する。
実行環境
- WSL2上のDockerで実行
- CUDA:10.2
- cudnn:7
- pytorch-lightning: 1.5.10
- torch:1.10.1
- tensorflow:2.6.2
- Nvidia driver:xxx
学習用Docker Imageを作成
CUDA10.2のDocker Imageを作成したら、以下のコマンドで学習用のDocker Imgageを作成する。
cd Segformer
docker build -t segformer_tensorboard:latest .
Dockerコンテナを起動
イメージを作成出来たら以下のコマンドでコンテナを起動します。Tensorboardで結果を確認したりするための引数を設定しています。
sudo docker run --gpus all --shm-size 32G -it --rm -p 10000:8888 -p 6006:6006 -u 0 --name segformer_tensorboard segformer_tensorboard:latest
学習
コンテナを起動出来たら以下のコマンドで学習を実行できます。
python3 segformer.py
学習時のloss配下のコマンドで可視化できます。
tensorboard --logdir="./lightning_logs/" --bind_all
http://localhost:6006
推論
推論は以下のコマンドで実行します。
python3 segformer_inference.py
検出結果
以下のようにいい感じに推論できていそうです。
DATALOADER:0 TEST RESULTS
{'test_loss': 0.06342098116874695,
'test_mean_accuracy': 0.9768195033576641,
'test_mean_iou': 0.9511022762808763}
Discussion