segnet网络预训练
SegNet是一个用于图像分割的卷积神经网络模型,一般情况下是使用已有的数据集进行训练。如果要进行预训练,可以考虑使用类似于图像分类任务中的预训练方法。可以使用ImageNet等大规模数据集进行预训练,然后在图像分割任务中微调模型。另外,也可以使用类似于自监督学习等方法进行预训练。但需要注意的是,预训练方法需要根据具体的任务需求进行选择和调整,不能一概而论。
pytorch的segnet训练自己的数据集
使用PyTorch实现SegNet对自定义数据集进行训练
准备工作环境
为了确保能够顺利运行代码,需安装必要的库并配置好开发环境。这通常涉及Python及其科学计算包numpy、图像处理库PIL以及最重要的深度学习框架PyTorch。
pip install numpy pillow torch torchvision
下载预训练模型或源码
可以从GitHub上获取已经编写好的SegNet网络结构作为起点[^3]:
git clone https://github.com/meetshah1995/pytorch-semseg.git
cd pytorch-semseg
此仓库不仅提供了完整的SegNet架构实现,还包含了其他几种流行的语义分割算法供参考。
构建自定义数据集类
对于特定应用的数据集来说,可能无法直接使用标准公开的数据集格式。因此需要继承torch.utils.data.Dataset
来创建适合项目需求的新类,在此类内部重写__len__()
方法返回样本总数目;通过覆写__getitem__(self, idx)
函数完成单个实例的读取操作,包括但不限于图片路径解析、标签映射转换等逻辑[^2]。
from PIL import Image
import os.path as osp
from glob import glob
class CustomDataset(torch.utils.data.Dataset):
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
# 假设数据集中每张图像是以其对应的索引命名的
image_files = sorted(glob(osp.join(root_dir,'images/*.png')))
label_files = sorted(glob(osp.join(root_dir,'labels/*.png')))
assert len(image_files)==len(label_files), "Image and Label count mismatch"
self.files = [{'image': img_path, 'label': lbl_path} for (img_path,lbl_path) in zip(image_files,label_files)]
def __len__(self):
return len(self.files)
def __getitem__(self, index):
datafile = self.files[index]
with open(datafile['image'], 'rb') as f:
image = Image.open(f).convert('RGB')
with open(datafile['label'], 'rb') as f:
label = Image.open(f).convert('L')
if self.transform is not None:
image, label = self.transform((image, label))
return {'X': image, 'Y': label}
创建DataLoader对象用于迭代访问批次化后的数据
有了上述定制化的数据集之后,下一步就是利用内置工具——DataLoader
—将其封装成易于使用的批处理形式,从而简化后续训练流程中的输入管理过程。
batch_size = 4
num_workers = multiprocessing.cpu_count()
train_dataset = CustomDataset('/path/to/trainset', transforms.Compose([
RandomCrop(size=(256, 256)),
ToTensor(),
]))
val_dataset = CustomDataset('/path/to/validationset', transforms.ToTensor())
train_loader = DataLoader(
dataset=train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers
)
val_loader = DataLoader(
dataset=val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers
)
定义损失函数与优化器
选择合适的评价指标和更新策略是提高模型性能的关键因素之一。这里采用交叉熵误差衡量预测值同真实标记之间的差异程度,并借助随机梯度下降法(SGD)调整权重参数以最小化这种差距[^1]。
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
开始训练循环
最后一步便是执行实际的学习任务了。一般情况下会重复多轮次遍历整个训练集合直至达到预期效果为止。期间还需定期评估当前状态下模型的表现情况以便及时发现潜在问题所在。
for epoch in range(num_epochs):
running_loss = 0.0
model.train()
for i_batch, sample_batched in enumerate(train_loader):
inputs = Variable(sample_batched['X'].cuda())
labels = Variable(sample_batched['Y'].squeeze().long().cuda())
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i_batch % log_interval == log_interval-1:
print('[%d/%d][%d/%d]\tLoss:%.4f' %(epoch+1,num_epochs,i_batch+1,len(train_loader),running_loss/log_interval))
running_loss = 0.0
validate_one_epoch(val_loader,model,criterion)
segnet跑通
成功运行 SegNet 模型实现图像分割
为了成功运行 SegNet 深度学习模型并实现图像分割,需遵循一系列特定的操作流程。这些操作不仅涉及环境搭建、数据准备,还包括模型的具体配置与训练过程。
环境安装与依赖库设置
首先,确保已安装 Python 和 PyTorch 环境。对于 SegNet 来说,推荐使用 Anaconda 虚拟环境管理工具来创建独立的工作空间,并通过 pip 或 conda 命令安装必要的包,如 torchvision, numpy, matplotlib 等[^1]。
conda create --name segnet python=3.9
conda activate segnet
pip install torch torchvision torchaudio
数据集准备
针对基建裂缝语义分割的任务,拥有一个高质量的数据集至关重要。该数据集中应包含大量标注过的基础设施裂缝图片及其对应的像素级标签图。按照标准做法,通常会将整个数据集划分为训练集、验证集和测试集三部分。
创建自定义 Dataset 类继承 torch.utils.data.Dataset
并重写其方法以适应具体应用场景下的读取逻辑:
from PIL import Image
import os
from torch.utils.data import Dataset
class CracksDataset(Dataset):
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.image_files = sorted(os.listdir(os.path.join(root_dir, 'images')))
self.mask_files = sorted(os.listdir(os.path.join(root_dir, 'masks')))
def __len__(self):
return len(self.image_files)
def __getitem__(self, idx):
img_path = os.path.join(self.root_dir, 'images', self.image_files[idx])
mask_path = os.path.join(self.root_dir, 'masks', self.mask_files[idx])
image = Image.open(img_path).convert('RGB')
mask = Image.open(mask_path)
if self.transform is not None:
transformed = self.transform(image=image, mask=mask)
image = transformed['image']
mask = transformed['mask']
return {'image': image, 'label': mask}
构建 SegNet 模型结构
SegNet 是一种经典的编码解码架构网络,它由 VGG16 提供预训练权重作为编码器部分的基础参数初始化方案;而解码阶段则采用上采样的方式逐步恢复原始分辨率特征映射。可以利用官方文档或者第三方开源项目中的实现版本快速构建起基础框架。
import torch.nn as nn
import torchvision.models.vgg as vgg
def make_layers(cfg, batch_norm=False):
layers = []
in_channels = 3
for v in cfg:
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
if batch_norm:
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
else:
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
return nn.Sequential(*layers)
cfgs = {
'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512],
}
class SegNet(nn.Module):
def __init__(self, num_classes=21):
super(SegNet, self).__init__()
# Encoder part using pretrained VGG16 weights
encoder = list(vgg.vgg16(pretrained=True).features.children())
self.encoders = [
nn.Sequential(*encoder[:5]),
nn.Sequential(*encoder[5:10]),
nn.Sequential(*encoder[10:17]),
nn.Sequential(*encoder[17:24]),
nn.Sequential(*encoder[24:])
]
# Decoder part with upsampling operations
decoder_cfg = cfgs['A'][:-1][::-1]
self.decoders = nn.ModuleList([
make_layers(decoder_cfg[i:i+2], True)[::-1]
for i in range(0, len(decoder_cfg)-1, 2)])
self.final_conv = nn.ConvTranspose2d(
64,
num_classes,
kernel_size=(2, 2),
stride=(2, 2))
def forward(self, x):
indices_list = []
for enc_layer in self.encoders:
before_pooling = x.clone()
x = enc_layer(x)
if isinstance(enc_layer[-1], nn.MaxPool2d):
_, indices = F.max_pool2d_with_indices(before_pooling, **enc_layer[-1].kwargs)
indices_list.append(indices)
for dec_layer, ind in zip(reversed(self.decoders), reversed(indices_list)):
x = F.max_unpool2d(x, ind, **dec_layer[-1].kwargs)
x = dec_layer(x)
output = self.final_conv(x)
return output
训练与评估
完成上述准备工作之后就可以着手编写训练脚本了。这一步骤涉及到损失函数的选择(交叉熵)、优化算法设定以及迭代更新策略等方面的内容。考虑到 GPU 加速的需求,建议尽可能多地利用 CUDA 设备加速计算效率。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SegNet(num_classes=len(CLASSES)).to(device)
criterion = nn.CrossEntropyLoss().to(device)
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
for epoch in range(NUM_EPOCHS):
model.train()
running_loss = 0.0
for data in dataloader_train:
inputs = data["image"].float().to(device)
labels = data["label"].long().squeeze(dim=-1).to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
avg_loss = running_loss / len(dataloader_train.dataset)
print(f'Epoch [{epoch}/{NUM_EPOCHS}], Loss: {avg_loss:.4f}')
相关推荐
















