
图像分割语义分割MobilenetV1-UNet网络:基于PyTorch框架全套项目,含模型、训练
、预测代码,直接下载即用
直接上干货!这次带来的MobileNetV1-UNet语义分割项目,绝对是懒人福音。整套代码从模型到训
练预测全打通,解压数据集就能跑。先看效果:在街景分割任务中,用1050Ti显卡训练两小时就能区分人行
道、车辆、建筑,预测单张图只要0.3秒。
核心代码都在`model.py`里。编码器用MobileNetV1的深度可分离卷积,参数只有原版UNet的1/3:
```python
class DepthwiseSeparable(nn.Module):
def __init__(self, in_ch, out_ch):
super().__init__()
self.depth_conv = nn.Conv2d(in_ch, in_ch, kernel_size=3,
padding=1, groups=in_ch)
self.point_conv = nn.Conv2d(in_ch, out_ch, kernel_size=1)
def forward(self, x):
return self.point_conv(self.depth_conv(x))
```
这个深度可分离卷积块是MobileNet的精髓——先做通道内卷积,再用1x1卷积混合通道。实测参数量
比标准卷积少8-9倍,且精度损失不大。
解码器部分保持UNet的跳跃连接结构,但上采样用了最近邻插值+卷积的方案:
```python
class DecodeBlock(nn.Module):
def __init__(self, in_ch, skip_ch, out_ch):
super().__init__()
self.up = nn.Upsample(scale_factor=2, mode='nearest')
self.conv = nn.Sequential(
nn.Conv2d(in_ch+skip_ch, out_ch, 3, padding=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True)
)