使用Pytorch搭建U-Net网络

popoier

已于 2023-10-07 21:00:12 修改

阅读量264

点赞数

文章标签： pytorch 人工智能 python

于 2023-10-07 20:53:15 首次发布

本文链接：https://blog.csdn.net/zoro_rn/article/details/133654939

版权

原因

github上关于Unet网络的实现不少，其中milesial实现了基于pytorch的，但是，在运行过程中，发现其代码训练很慢，而且特别占内存，在显存为12G的3060上的batch_szie也只能为2。故另寻其方法，好在b站博主霹雳吧啦Wz实现了pytorch的简化版本，这里我推荐一下这位博主，很适合初学者。

搭建U-Net

根据其在github上的readme，搭建好环境，只需要改变my_dataset.py文件即可运行自己的数据集。
my_dataset.py更改如下：

import os
from PIL import Image
import numpy as np
from torch.utils.data import Dataset


class DriveDataset(Dataset):
    def __init__(self, root: str, train: bool, transforms=None):
        super(DriveDataset, self).__init__()
        self.flag = "training" if train else "test"
        data_root = os.path.join(root, "DRIVE", self.flag)
        assert os.path.exists(data_root), f"path '{data_root}' does not exists."
        self.transforms = transforms
        img_names = [i for i in os.listdir(os.path.join(data_root, "images")) if i.endswith(".jpg")]
        self.img_list = [os.path.join(data_root, "images", i) for i in img_names]
        mask_names = [i for i in os.listdir(os.path.join(data_root, "mask")) if i.endswith(".png")]
        self.mask_list = [os.path.join(data_root, "mask", i) for i in mask_names]

    def __getitem__(self, idx):
        img = Image.open(self.img_list[idx]).convert('RGB')
        mask = Image.open(self.mask_list[idx])

        if self.transforms is not None:
            img, mask = self.transforms(img, mask)

        return img, mask

    def __len__(self):
        return len(self.img_list)

    @staticmethod
    def collate_fn(batch):
        images, targets = list(zip(*batch))
        batched_imgs = cat_list(images, fill_value=0)
        batched_targets = cat_list(targets, fill_value=255)
        return batched_imgs, batched_targets


def cat_list(images, fill_value=0):
    max_size = tuple(max(s) for s in zip(*[img.shape for img in images]))
    batch_shape = (len(images),) + max_size
    batched_imgs = images[0].new(*batch_shape).fill_(fill_value)
    for img, pad_img in zip(images, batched_imgs):
        pad_img[..., :img.shape[-2], :img.shape[-1]].copy_(img)
    return batched_imgs