Pytorch UNet

U-NET：pytorch的語義分割

從高清圖像中的Kaggle的Carvana Image掩蓋挑戰中，在Pytorch中定制了U-NET實現。

• 快速開始
  • 沒有Docker
  • 與Docker
• 描述
• 用法
  • Docker
  • 訓練
  • 預言
• 重量和偏見
• 預驗證的模型
• 數據

快速開始

沒有Docker

1. 安裝CUDA

2. 安裝Pytorch 1.13或更高版本

3. 安裝依賴項

pip install -r requirements.txt

4. 下載數據並運行培訓：

bash scripts/download_data.sh
python train.py --amp

與Docker

1. 安裝Docker 19.03或更高版本：

curl https://get.docker.com | sh && sudo systemctl --now enable docker

2. 安裝NVIDIA容器工具包：

distribution= $( . /etc/os-release ; echo $ID$VERSION_ID ) 
   && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - 
   && curl -s -L https://nvidia.github.io/nvidia-docker/ $distribution /nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-docker2
sudo systemctl restart docker

3. 下載並運行圖像：

sudo docker run --rm --shm-size=8g --ulimit memlock=-1 --gpus all -it milesial/unet

4. 下載數據並運行培訓：

bash scripts/download_data.sh
python train.py --amp

描述

該模型從刮擦訓練了5K圖像，並在超過100K的測試圖像上得分為0.988423。

它可以輕鬆用於多類分割，肖像細分，醫學細分，...

用法

注意：使用Python 3.6或更新

Docker

包含代碼和依賴項的Docker映像可在Dockerhub上獲得。您可以使用（Docker> = 19.03）下載並跳入容器中：

 docker run -it --rm --shm-size=8g --ulimit memlock=-1 --gpus all milesial/unet

訓練

> python train.py -h
usage: train.py [-h] [--epochs E] [--batch-size B] [--learning-rate LR]
                [--load LOAD] [--scale SCALE] [--validation VAL] [--amp]

Train the UNet on images and target masks

optional arguments:
  -h, --help            show this help message and exit
  --epochs E, -e E      Number of epochs
  --batch-size B, -b B  Batch size
  --learning-rate LR, -l LR
                        Learning rate
  --load LOAD, -f LOAD  Load model from a .pth file
  --scale SCALE, -s SCALE
                        Downscaling factor of the images
  --validation VAL, -v VAL
                        Percent of the data that is used as validation (0-100)
  --amp                 Use mixed precision

默認情況下， scale為0.5，因此，如果您希望獲得更好的結果（但使用更多內存），請將其設置為1。

--amp標誌也可以使用自動混合精度。混合精度使模型可以使用FP16算術使用更少的內存，並在最近的GPU上更快。建議啟用放大器。

預言

訓練模型並將其保存到MODEL.pth之後，您可以通過CLI輕鬆測試圖像上的輸出掩碼。

預測單個圖像並保存它：

python predict.py -i image.jpg -o output.jpg

預測多個圖像並在不保存它們的情況下向它們顯示：

python predict.py -i image1.jpg image2.jpg --viz --no-save

> python predict.py -h
usage: predict.py [-h] [--model FILE] --input INPUT [INPUT ...] 
                  [--output INPUT [INPUT ...]] [--viz] [--no-save]
                  [--mask-threshold MASK_THRESHOLD] [--scale SCALE]

Predict masks from input images

optional arguments:
  -h, --help            show this help message and exit
  --model FILE, -m FILE
                        Specify the file in which the model is stored
  --input INPUT [INPUT ...], -i INPUT [INPUT ...]
                        Filenames of input images
  --output INPUT [INPUT ...], -o INPUT [INPUT ...]
                        Filenames of output images
  --viz, -v             Visualize the images as they are processed
  --no-save, -n         Do not save the output masks
  --mask-threshold MASK_THRESHOLD, -t MASK_THRESHOLD
                        Minimum probability value to consider a mask pixel white
  --scale SCALE, -s SCALE
                        Scale factor for the input images

您可以指定要與--model MODEL.pth一起使用的模型文件。

重量和偏見

訓練進度可以使用權重和偏見實時可視化。損耗曲線，驗證曲線，權重和梯度直方圖以及預測的面具被記錄到平台。

在啟動培訓時，將在控制台中打印鏈接。單擊它進入儀表板。如果您有一個現有的W＆B帳戶，則可以通過設置WANDB_API_KEY環境變量來鏈接它。如果沒有，它將創建一個匿名運行，該運行將在7天后自動刪除。

預驗證的模型

Carvana數據集可用審慎的型號。它也可以從Torch.hub上加載：

 net = torch . hub . load ( 'milesial/Pytorch-UNet' , 'unet_carvana' , pretrained = True , scale = 0.5 )

可用量表為0.5和1.0。

數據

Carvana數據可在Kaggle網站上找到。

您也可以使用助手腳本下載它：

 bash scripts/download_data.sh

輸入圖像和目標掩碼應分別在data/imgs和data/masks文件夾中（請注意，由於貪婪的數據加載程序， imgs和masks文件夾不應包含任何子文件夾或任何其他文件）。對於carvana，圖像是RGB，面具是黑色和白色。

只要確保在utils/data_loading.py中正確加載它，就可以使用自己的數據集。

---

Olaf Ronneberger，Philipp Fischer，Thomas Brox的原始論文：

U-NET：生物醫學圖像分割的捲積網絡