import segmentation_models_pytorchでエラー

環境

Windows 10 Pro
GPU: NVIDIA RTX 3090
CUDA: 11.0
Python version: 3.8.8
Pytorch version: 1.9 (記事中で 1.7.1に変更)

エラー

UNet++を実装する際、

examin.py

import segmentation_models_pytorch

を実行したところ、

ImportError: cannot import name 'container_abcs' from 'torch._six'

とエラーが発生した。

対処法

いろいろ調べてみると、（参考→https://github.com/NVIDIA/apex/pull/1049）
Pytorchのversionがよくなかったらしい。
version==1.8よりも前のversionでしか_six.pyからcontainer_abcsをインポートできないらしい。
現在のversionを調べた。

examin2.py

import torch
print(torch.__version__) # 1.9であった

これだ。versionを1.7.1に下げる。
まずは、現在のPytorchとsegmentation_models_pytorchをpip uninstallする。
https://pytorch.org/
のpreviousにアクセスして、自分のGPUやCUDAに合わせてpytorchを選択し、pip installで再インストールする。（もちろん、segmentation_models_pytorchも再インストール。）

examin3.py

import torch
print(torch.__version__) # 1.7.1になってる！
print(torch.cuda.is_available) # True! OK!!
import segmentation_models_pytorch # OK! 2021/6/29 追加

ということで、何とかなりました！

コメント

初心者なので、versionという初歩的なところでこけてしまいました。
よかったら、参考にしてみてください。

何か、コメント等ありましたら、お願いします！！