全卷积网络（FCN）实战：使用FCN实现语义分割(2)

root=root.replace(‘\’,‘/’)

assert os.path.exists(root), “path ‘{}’ does not exist.”.format(root)

image_dir = os.path.join(root, ‘JPEGImages’)

mask_dir = os.path.join(root, ‘SegmentationClass’)

txt_path = os.path.join(root, “ImageSets”, “Segmentation”, txt_name)

txt_path=txt_path.replace(‘\’,‘/’)

assert os.path.exists(txt_path), “file ‘{}’ does not exist.”.format(txt_path)

with open(os.path.join(txt_path), “r”) as f:

file_names = [x.strip() for x in f.readlines() if len(x.strip()) > 0]

self.images = [os.path.join(image_dir, x + “.jpg”) for x in file_names]

self.masks = [os.path.join(mask_dir, x + “.png”) for x in file_names]

assert (len(self.images) == len(self.masks))

self.transforms = transforms

导入需要的包。

定义VOC数据集读取类VOCSegmentation。在init方法中，核心是读取image列表和mask列表。

def getitem(self, index):

img = Image.open(self.images[index]).convert(‘RGB’)

target = Image.open(self.masks[index])

if self.transforms is not None:

img, target = self.transforms(img, target)

return img, target

__getitem__方法是获取单张图片和图片对应的mask，然后对其做数据增强。

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

collate_fn方法是对一个batch中数据调用cat_list做数据对齐。

在train.py中torch.utils.data.DataLoader调用

train_loader = torch.utils.data.DataLoader(train_dataset,

batch_size=batch_size,

num_workers=num_workers,

shuffle=True,

pin_memory=True,

collate_fn=train_dataset.collate_fn)

val_loader = torch.utils.data.DataLoader(val_dataset,

batch_size=1,

num_workers=num_workers,

pin_memory=True,

collate_fn=val_dataset.collate_fn)

训练

=============================================================

重要参数

---

打开train.py,我们先认识一下重要的参数：

def parse_args():

import argparse

parser = argparse.ArgumentParser(description=“pytorch fcn training”)

数据集的根目录(VOCdevkit)所在的文件夹

parser.add_argument(“–data-path”, default=“data/”, help=“VOCdevkit root”)

parser.add_argument(“–num-classes”, default=20, type=int)

parser.add_argument(“–aux”, default=True, type=bool, help=“auxilier loss”)

parser.add_argument(“–device”, default=“cuda”, help=“training device”)

parser.add_argument(“-b”, “–batch-size”, default=32, type=int)

parser.add_argument(“–epochs”, default=30, type=int, metavar=“N”,

help=“number of total epochs to train”)

parser.add_argument(‘–lr’, default=0.0001, type=float, help=‘initial learning rate’)

parser.add_argument(‘–momentum’, default=0.9, type=float, metavar=‘M’,

help=‘momentum’)

parser.add_argument(‘–wd’, ‘–weight-decay’, default=1e-4, type=float,

metavar=‘W’, help=‘weight decay (default: 1e-4)’,

dest=‘weight_decay’)

parser.add_argument(‘–print-freq’, default=10, type=int, help=‘print frequency’)

parser.add_argument(‘–resume’, default=‘’, help=‘resume from checkpoint’)

parser.add_argument(‘–start-epoch’, default=0, type=int, metavar=‘N’,

help=‘start epoch’)

是否使用混合精度训练

parser.add_argument(“–amp”, default=False, type=bool,

help=“Use torch.cuda.amp for mixed precision training”)

args = parser.parse_args()

return args

data-path：定义数据集的根目录(VOCdevkit)所在的文件夹

num-classes：检测目标类别数(不包含背景)。

aux：是否使用aux_classifier。

device：使用cpu还是gpu训练，默认是cuda。

batch-size：BatchSize设置。

epochs：epoch的个数。

lr：学习率。

resume：继续训练时候，选择用的模型。

start-epoch：起始的epoch，针对再次训练时，可以不需要从0开始。

amp：是否使用torch的自动混合精度训练。

数据增强

---

增强调用transforms.py中的方法。

训练集的增强如下：

class SegmentationPresetTrain:

def init(self, base_size, crop_size, hflip_prob=0.5, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):

随机Resize的最小尺寸

min_size = int(0.5 * base_size)

随机Resize的最大尺寸

max_size = int(2.0 * base_size)

随机Resize增强。

trans = [T.RandomResize(min_size, max_size)]

if hflip_prob > 0:

#随机水平翻转

trans.append(T.RandomHorizontalFlip(hflip_prob))

trans.extend([

#随机裁剪

T.RandomCrop(crop_size),

T.ToTensor(),

T.Normalize(mean=mean, std=std),

])

self.transforms = T.Compose(trans)

def call(self, img, target):

return self.transforms(img, target)

训练集增强，包括随机Resize、随机水平翻转、随即裁剪。

验证集增强：

class SegmentationPresetEval:

def init(self, base_size, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):

self.transforms = T.Compose([

T.RandomResize(base_size, base_size),

T.ToTensor(),

T.Normalize(mean=mean, std=std),

])

def call(self, img, target):

return self.transforms(img, target)

验证集的增强比较简单，只有随机Resize。

Main方法

---

对Main方法，我做了一些修改，修改的代码如下：

#定义模型，并加载预训练

model = fcn_resnet50(pretrained=True)

默认classes是21,如果不是21，则要修改类别。

if num_classes != 21:

model.classifier[4] = torch.nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1))

model.aux_classifier[4] = torch.nn.Conv2d(256, num_classes, kernel_size=(1, 1), stride=(1, 1))

print(model)

model.to(device)

如果有多张显卡，则使用多张显卡

if torch.cuda.device_count() > 1:

print(“Let’s use”, torch.cuda.device_count(), “GPUs!”)

model = torch.nn.DataParallel(model)

模型，我改为pytorch官方的模型了，如果能使用官方的模型尽量使用官方的模型。

默认类别是21，如果不是21，则要修改类别。

检测系统中是否有多张卡，如果有多张卡则使用多张卡不能浪费资源。

如果不想使用所有的卡，而是指定其中的几张卡，可以使用：

os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘0,1’

也可以在DataParallel方法中设定：

model = torch.nn.DataParallel(model,device_ids=[0,1])

如果使用了多显卡，再使用模型的参数就需要改为model.module.xxx，例如：

params = [p for p in model.module.aux_classifier.parameters() if p.requires_grad]

params_to_optimize.append({“params”: params, “lr”: args.lr * 10})

上面的都完成了就可以开始训练了，如下图：

测试

=====================================================================================================================================================================

在开始测试之前，我们还要获取到调色板，新建脚本get_palette.py，代码如下：

import json

import numpy as np

from PIL import Image

读取mask标签

target = Image.open(“./2007_001288.png”)

获取调色板

palette = target.getpalette()

palette = np.reshape(palette, (-1, 3)).tolist()

print(palette)

转换成字典子形式

pd = dict((i, color) for i, color in enumerate(palette))

json_str = json.dumps(pd)

with open(“palette.json”, “w”) as f:

f.write(json_str)

选取一张mask，然后使用getpalette方法获取，然后将其转为字典的格式保存。

接下来，开始预测部分，新建predict.py,插入以下代码：

import os

import time

import json

import torch

from torchvision import transforms

import numpy as np

from PIL import Image

from torchvision.models.segmentation import fcn_resnet50

导入程序需要的包文件，然在mian方法中：

def main():

aux = False # inference time not need aux_classifier

classes = 20

weights_path = “./save_weights/model_5.pth”

img_path = “./2007_000123.jpg”

palette_path = “./palette.json”

assert os.path.exists(weights_path), f"weights {weights_path} not found."

assert os.path.exists(img_path), f"image {img_path} not found."

assert os.path.exists(palette_path), f"palette {palette_path} not found."

with open(palette_path, “rb”) as f:

pallette_dict = json.load(f)

pallette = []

for v in pallette_dict.values():

pallette += v

• 定义是否需要aux_classifier，预测不需要aux_classifier，所以设置为False。

• 设置类别为20，不包括背景。

• 定义权重的路径。

• 定义调色板的路径。

• 读去调色板。

接下来，是加载模型，单显卡训练出来的模型和多显卡训练出来的模型加载有区别，我们先看单显卡训练出来的模型如何加载。

model = fcn_resnet50(num_classes=classes+1)

print(model)

单显卡训练出来的模型，加载

delete weights about aux_classifier

weights_dict = torch.load(weights_path, map_location=‘cpu’)[‘model’]

for k in list(weights_dict.keys()):

if “aux_classifier” in k:

del weights_dict[k]

load weights

model.load_state_dict(weights_dict)

model.to(device)

定义模型fcn_resnet50，num_classes设置为类别+1（背景）

加载训练好的模型，并将aux_classifier删除。

然后加载权重。

再看多显卡的模型如何加载

create model

model = fcn_resnet50(num_classes=classes+1)

model = torch.nn.DataParallel(model)

delete weights about aux_classifier

weights_dict = torch.load(weights_path, map_location=‘cpu’)[‘model’]

print(weights_dict)

for k in list(weights_dict.keys()):

if “aux_classifier” in k:

del weights_dict[k]

load weights

model.load_state_dict(weights_dict)

model=model.module

model.to(device)

定义模型fcn_resnet50，num_classes设置为类别+1（背景），将模型放入DataParallel类中。

加载训练好的模型，并将aux_classifier删除。

加载权重。

执行torch.nn.DataParallel(model)时，model被放在了model.module，所以model.module才真正需要的模型。所以我们在这里将model.module赋值给model。

接下来是图像数据的处理

load image

original_img = Image.open(img_path)

from pil image to tensor and normalize

data_transform = transforms.Compose([transforms.Resize(520),

transforms.ToTensor(),

transforms.Normalize(mean=(0.485, 0.456, 0.406),

（1）Python所有方向的学习路线（新版）

这是我花了几天的时间去把Python所有方向的技术点做的整理，形成各个领域的知识点汇总，它的用处就在于，你可以按照上面的知识点去找对应的学习资源，保证自己学得较为全面。

最近我才对这些路线做了一下新的更新，知识体系更全面了。

（2）Python学习视频

包含了Python入门、爬虫、数据分析和web开发的学习视频，总共100多个，虽然没有那么全面，但是对于入门来说是没问题的，学完这些之后，你可以按照我上面的学习路线去网上找其他的知识资源进行进阶。

（3）100多个练手项目

我们在看视频学习的时候，不能光动眼动脑不动手，比较科学的学习方法是在理解之后运用它们，这时候练手项目就很适合了，只是里面的项目比较多，水平也是参差不齐，大家可以挑自己能做的项目去练练。

网上学习资料一大堆，但如果学到的知识不成体系，遇到问题时只是浅尝辄止，不再深入研究，那么很难做到真正的技术提升。

需要这份系统化学习资料的朋友，可以戳这里无偿获取

一个人可以走的很快，但一群人才能走的更远！不论你是正从事IT行业的老鸟或是对IT行业感兴趣的新人，都欢迎加入我们的的圈子（技术交流、学习资源、职场吐槽、大厂内推、面试辅导），让我们一起学习成长！