// the find
bubbliiiing/unet-pytorch
这是一个unet-pytorch的源码,可以训练自己的模型
A PyTorch implementation of U-Net for semantic segmentation, targeting both medical imaging (the architecture's original purpose) and general datasets like VOC. It supports VGG and ResNet50 backbones, includes mIOU evaluation, and ships with a small example medical dataset. Aimed at someone who wants a working U-Net baseline without building one from scratch.
The author is honest that U-Net doesn't suit VOC-style datasets and is better for medical/sparse-feature tasks — that caveat is right in the README. Dual backbone support (VGG/ResNet50) with pretrained weights is genuinely useful for quick experiments. The repo includes a tiny bundled medical dataset so you can actually run it end-to-end without hunting for data. Training loop supports cosine and step LR schedules plus Adam/SGD with batch-size-adaptive LR, which is more than most educational U-Net repos bother with.
Pinned to torch==1.2.0 from 2019 — that's not a typo, it's just abandoned. Modern PyTorch has changed enough that you'll spend time debugging compatibility before training a single epoch. Documentation is entirely in Chinese with no English translation, which closes the door on most of the potential audience. The medical dataset is 30 images, which is fine for smoke-testing but meaningless for actual training or benchmarking. Last commit was May 2024 but the architecture itself hasn't evolved — no attention gates, no transformer bottleneck, nothing from the last four years of U-Net variants that actually matter in practice.