共计 12307 个字符，预计需要花费 31 分钟才能阅读完成。
摘要：U2Net 是一个优良的显著性指标检测算法，由 Qin Xuebin 等人发表在 Pattern Recognition 2020 期刊[Arxiv]。U2Net 名称的起源在于其网络结构由两层嵌套的 Unet 构造，能够在不须要预训练骨干网络的状况下从零开始训练，领有优异的体现。
本文分享自华为云社区《ModelArts Notebook 疾速开源我的项目实战 — U2Net》，作者：shpity。
一、U2Net 介绍
U2Net 是一个优良的显著性指标检测算法，由 Qin Xuebin 等人发表在 Pattern Recognition 2020 期刊[Arxiv]。U2Net 名称的起源在于其网络结构由两层嵌套的 Unet 构造，能够在不须要预训练骨干网络的状况下从零开始训练，领有优异的体现。其网络结构如图 1 所示。
图 1. U2Net 的主体框架是一个相似于 U -Net 的编解码构造，然而每一个 block 替换为新提出的残差 U -block 模块
我的项目开源地址：https://github.com/xuebinqin/…
二、创立 Notebook 开发环境
1. 进入 ModelArts 控制台
2. 抉择开发环境 -> Notebook -> 创立
3. 创立 Notebook
3.1 能够抉择和工作相干的名称，方便管理；
3.2 为了缩小不必要的资源耗费，倡议开启主动进行；
3.3 U2Net 所需的运行环境在公共镜像中曾经蕴含，能够抉择 pytorch1.4-cuda10.1-cudnn7-ubuntu18.04；
3.4 倡议抉择 GPU 类型，不便模型疾速训练；
3.5 抉择立刻创立 -> 提交，期待 notebook 创立实现后关上 Notebook。
4. 导入开源我的项目源码（git/ 手动上传）
4.1 在 Terminal 应用 git 克隆近程仓库
cd work # 留神：只有 /home/ma-user/work 目录及其子目录下的文件在 Notebook 实例敞开后会保留
git clone https://github.com/xuebinqin/U-2-Net.git4.2 如果 git 速度较慢也能够从本地上传代码，间接将压缩包拖到左侧文件目录栏或者采纳 OBS 上传。
三、数据筹备
1. 下载训练数据 APDrawing dataset
应用 Wget 间接下载到 Notebook，也可下载本地后再拖拽到 Notebook 中。
wget https://cg.cs.tsinghua.edu.cn/people/~Yongjin/APDrawingDB.zip
unzip APDrawingDB.zip注：如果数据集较大（>5GB）须要下载到其它目录（实例进行后会被删除），倡议寄存在 OBS 中，须要的时候随时拉取。
# 从 OBS 中拉取代码到指定目录
sh-4.4$ source /home/ma-user/anaconda3/bin/activate PyTorch-1.4
sh-4.4$ python
>>> mox.file.copy_parallel('obs://bucket-xxxx/APDrawingDB', '/home/ma-user/work/APDrawingDB')2. 切分训练数据
数据集中./APDrawingDB/data/train 中蕴含了 420 张训练图片，分辨率为 512*1024，左侧为输出图像，右侧为对应的 ground truth。咱们须要将大图从两头切分为两个子图。
2.1 在 Notebook 开发环境中新建一个 Pytorch-1.4 的 jupyter Notebook 文件，名称能够为 split.ipynb，脚本将会在./APDrawingDB/data/train/split 目录下生成 840 张子图，其中原始图像以.jpg 结尾，gt 图像以.png 结尾，不便后续训练代码读取【test 文件夹切分步骤同理】。
from PIL import Image
import os
train_img_dir = os.path.join("./APDrawingDB/data/train")
img_list = os.listdir(train_img_dir)
for image in img_list:
img_path = os.path.join(train_img_dir, image)
if not os.path.isdir(img_path):
img = Image.open(img_path)
#print(img.size)
save_img_dir = os.path.join(train_img_dir, 'split_train')
if not os.path.exists(save_img_dir):
os.mkdir(save_img_dir)
save_img_path = os.path.join(save_img_dir, image)
cropped_left = img.crop((0, 0, 512, 512)) # (left, upper, right, lower)
cropped_right = img.crop((512, 0, 1024, 512)) # (left, upper, right, lower)
cropped_left.save(save_img_path[:-3] + 'jpg')
cropped_right.save(save_img_path)
test_img_dir = os.path.join("./APDrawingDB/data/test")
img_list = os.listdir(test_img_dir)
for image in img_list:
img_path = os.path.join(test_img_dir, image)
if not os.path.isdir(img_path):
img = Image.open(img_path)
#print(img.size)
save_img_dir = os.path.join(test_img_dir, 'split')
if not os.path.exists(save_img_dir):
os.mkdir(save_img_dir)
save_img_path = os.path.join(save_img_dir, image)
cropped_left = img.crop((0, 0, 512, 512)) # (left, upper, right, lower)
cropped_right = img.crop((512, 0, 1024, 512)) # (left, upper, right, lower)
cropped_left.save(save_img_path[:-3] + 'jpg')3. 将切分好的数据依照如下层级构造整顿出训练和测试所需的 datasets 文件夹
datasets/
├── test (70 张切分图片，只蕴含原图)
└── train (840 张切分图片，蕴含 420 张原图及对应的 gt)
注：能够将切分好的数据集保留到 OBS 目录中，缩小./work 的磁盘空间占用。
4. 残缺的 U -2-Net 我的项目构造如下所示：
U-2-Net/
├── .git
├── LICENSE
├── README.md
├── pycache
├── clipping_camera.jpg
├── data_loader.py
├── datasets
├── figures
├── gradio
├── model
├── requirements.txt
├── saved_models
├── setup_model_weights.py
├── test_data
├── u2net_human_seg_test.py
├── u2net_portrait_demo.py
├── u2net_portrait_test.py
├── u2net_test.py
└── u2net_train.py
四、训练
1. 官网提供的训练代码中数据的门路和咱们的 datasets 有些区别，须要对训练脚本进行一些批改，倡议应用 jupyter notebook 不便排除谬误
新建一个 Pytorch-1.4 的 jupyter Notebook 文件，名称能够为 train.ipynb
import moxing as mox
# 如果须要从 OBS 拷贝切分好的训练数据
#mox.file.copy_parallel('obs://bucket-test-xxxx', '/home/ma-user/work/U-2-Net/datasets')
INFO:root:Using MoXing-v1.17.3-43fbf97f
INFO:root:Using OBS-Python-SDK-3.20.7
import os
import torch
import torchvision
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import torch.optim as optim
import torchvision.transforms as standard_transforms
import numpy as np
import glob
import os
from data_loader import Rescale
from data_loader import RescaleT
from data_loader import RandomCrop
from data_loader import ToTensor
from data_loader import ToTensorLab
from data_loader import SalObjDataset
from model import U2NET
from model import U2NETP
/home/ma-user/anaconda3/envs/PyTorch-1.4/lib/python3.7/site-packages/skimage/io/manage_plugins.py:23: UserWarning: Your installed pillow version is < 7.1.0. Several security issues (CVE-2020-11538, CVE-2020-10379, CVE-2020-10994, CVE-2020-10177) have been fixed in pillow 7.1.0 or higher. We recommend to upgrade this library.
from .collection import imread_collection_wrapper
bce_loss = nn.BCELoss(size_average=True)
/home/ma-user/anaconda3/envs/PyTorch-1.4/lib/python3.7/site-packages/torch/nn/_reduction.py:43: UserWarning: size_average and reduce args will be deprecated, please use reduction='mean' instead.
warnings.warn(warning.format(ret))
def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v):
loss0 = bce_loss(d0,labels_v)
loss1 = bce_loss(d1,labels_v)
loss2 = bce_loss(d2,labels_v)
loss3 = bce_loss(d3,labels_v)
loss4 = bce_loss(d4,labels_v)
loss5 = bce_loss(d5,labels_v)
loss6 = bce_loss(d6,labels_v)
loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6
print("l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n"%(loss0.data.item(),loss1.data.item(),loss2.data.item(),loss3.data.item(),loss4.data.item(),loss5.data.item(),loss6.data.item()))
return loss0, loss
model_name = 'u2net' #'u2netp'
data_dir = os.path.join(os.getcwd(), 'datasets', 'train' + os.sep)
# tra_image_dir = os.path.join('DUTS', 'DUTS-TR', 'DUTS-TR', 'im_aug' + os.sep)
# tra_label_dir = os.path.join('DUTS', 'DUTS-TR', 'DUTS-TR', 'gt_aug' + os.sep)
image_ext = '.jpg'
label_ext = '.png'
model_dir = os.path.join(os.getcwd(), 'saved_models', model_name + os.sep)
epoch_num = 100000
batch_size_train = 24
batch_size_val = 1
train_num = 0
val_num = 0
tra_img_name_list = glob.glob(data_dir + '*' + image_ext)
tra_lbl_name_list = []
for img_path in tra_img_name_list:
img_name = img_path.split(os.sep)[-1]
aaa = img_name.split(".")
bbb = aaa[0:-1]
imidx = bbb[0]
for i in range(1,len(bbb)):
imidx = imidx + "." + bbb[i]
tra_lbl_name_list.append(data_dir + imidx + label_ext)
print("---")
print("train images:", len(tra_img_name_list))
print("train labels:", len(tra_lbl_name_list))
print("---")
train_num = len(tra_img_name_list)
---
train images: 420
train labels: 420
---
salobj_dataset = SalObjDataset(
img_name_list=tra_img_name_list,
lbl_name_list=tra_lbl_name_list,
transform=transforms.Compose([RescaleT(320),
RandomCrop(288),
ToTensorLab(flag=0)]))
salobj_dataloader = DataLoader(salobj_dataset, batch_size=batch_size_train, shuffle=True, num_workers=1)
# ------- 3. define model --------
# define the net
if(model_name=='u2net'):
net = U2NET(3, 1)
elif(model_name=='u2netp'):
net = U2NETP(3,1)
if torch.cuda.is_available():
net.cuda()
# ------- 4. define optimizer --------
print("---define optimizer...")
optimizer = optim.Adam(net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
---define optimizer...
# ------- 5. training process --------
print("---start training...")
ite_num = 0
running_loss = 0.0
running_tar_loss = 0.0
ite_num4val = 0
save_frq = 2000 # save the model every 2000 iterations
---start training...
for epoch in range(0, epoch_num):
net.train()
for i, data in enumerate(salobj_dataloader):
ite_num = ite_num + 1
ite_num4val = ite_num4val + 1
inputs, labels = data['image'], data['label']
inputs = inputs.type(torch.FloatTensor)
labels = labels.type(torch.FloatTensor)
# wrap them in Variable
if torch.cuda.is_available():
inputs_v, labels_v = Variable(inputs.cuda(), requires_grad=False), Variable(labels.cuda(),
requires_grad=False)
else:
inputs_v, labels_v = Variable(inputs, requires_grad=False), Variable(labels, requires_grad=False)
# y zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
d0, d1, d2, d3, d4, d5, d6 = net(inputs_v)
loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v)
loss.backward()
optimizer.step()
# # print statistics
running_loss += loss.data.item()
running_tar_loss += loss2.data.item()
# del temporary outputs and loss
del d0, d1, d2, d3, d4, d5, d6, loss2, loss
print("[epoch: %3d/%3d, batch: %5d/%5d, ite: %d] train loss: %3f, tar: %3f" % (epoch + 1, epoch_num, (i + 1) * batch_size_train, train_num, ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val))
if ite_num % save_frq == 0:
model_weight = model_dir + model_name+"_bce_itr_%d_train_%3f_tar_%3f.pth" % (ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val)
torch.save(net.state_dict(), model_weight)
mox.file.copy_parallel(model_weight, 'obs://bucket-xxxx/output/model_save/' + model_weight.split('/')[-1])
running_loss = 0.0
running_tar_loss = 0.0
net.train() # resume train
ite_num4val = 0
l0: 0.167562, l1: 0.153742, l2: 0.156246, l3: 0.163096, l4: 0.176632, l5: 0.197176, l6: 0.247590
[epoch: 1/100000, batch: 24/ 420, ite: 500] train loss: 1.189413, tar: 0.159183
l0: 0.188048, l1: 0.179041, l2: 0.180086, l3: 0.187904, l4: 0.198345, l5: 0.218509, l6: 0.269199
[epoch: 1/100000, batch: 48/ 420, ite: 501] train loss: 1.266652, tar: 0.168805
l0: 0.192491, l1: 0.187615, l2: 0.188043, l3: 0.197142, l4: 0.203571, l5: 0.222019, l6: 0.261745
[epoch: 1/100000, batch: 72/ 420, ite: 502] train loss: 1.313146, tar: 0.174727
l0: 0.169403, l1: 0.155883, l2: 0.157974, l3: 0.164012, l4: 0.175975, l5: 0.195938, l6: 0.244896
[epoch: 1/100000, batch: 96/ 420, ite: 503] train loss: 1.303333, tar: 0.173662
l0: 0.171904, l1: 0.157170, l2: 0.156688, l3: 0.162020, l4: 0.175565, l5: 0.200576, l6: 0.258133
[epoch: 1/100000, batch: 120/ 420, ite: 504] train loss: 1.299787, tar: 0.173369
l0: 0.177398, l1: 0.166131, l2: 0.169089, l3: 0.176976, l4: 0.187039, l5: 0.205449, l6: 0.248036五、测试
新建一个 Pytorch-1.4 的 jupyter Notebook 文件，名称能够为 test.ipynb
import moxing as mox
# 拷贝数据
mox.file.copy_parallel('obs://bucket-xxxx/output/model_save/u2net.pth', '/home/ma-user/work/U-2-Net/saved_models/u2net/u2net.pth')
import os
import sys
from skimage import io, transform
import torch
import torchvision
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms#, utils
# import torch.optim as optim
import numpy as np
from PIL import Image
import glob
from data_loader import RescaleT
from data_loader import ToTensor
from data_loader import ToTensorLab
from data_loader import SalObjDataset
from model import U2NET # full size version 173.6 MB
from model import U2NETP # small version u2net 4.7 MB
%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib.pyplot import imshow
# normalize the predicted SOD probability map
def normPRED(d):
ma = torch.max(d)
mi = torch.min(d)
dn = (d-mi)/(ma-mi)
return dn
def save_output(image_name,pred,d_dir, show=False):
predict = pred
predict = predict.squeeze()
predict_np = predict.cpu().data.numpy()
im = Image.fromarray(predict_np*255).convert('RGB')
img_name = image_name.split(os.sep)[-1]
image = io.imread(image_name)
imo = im.resize((image.shape[1],image.shape[0]),resample=Image.BILINEAR)
pb_np = np.array(imo)
if show:
show_on_notebook(image, im)
aaa = img_name.split(".")
bbb = aaa[0:-1]
imidx = bbb[0]
for i in range(1,len(bbb)):
imidx = imidx + "." + bbb[i]
imo.save(d_dir+imidx+'.png')
return im
def show_on_notebook(image_original, pred): #此函数能够在 notebook 中展现模型的预测成果
plt.subplot(1,2,1)
imshow(np.array(image_original))
plt.subplot(1,2,2)
imshow(np.array(pred))
# --------- 1. get image path and name ---------
model_name='u2net'#u2netp
image_dir = os.path.join(os.getcwd(), 'datasets', 'test') #留神这里的 test_data/original 寄存的是 datasets/test 中的原始图片，不蕴含 gt
prediction_dir = os.path.join(os.getcwd(), 'output', model_name + '_results' + os.sep)
model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth')
img_name_list = glob.glob(os.path.join(os.getcwd(), 'datasets/test/*.jpg'))
# print(img_name_list)
# --------- 2. dataloader ---------
#1. dataloader
test_salobj_dataset = SalObjDataset(img_name_list = img_name_list,
lbl_name_list = [],
transform=transforms.Compose([RescaleT(320),
ToTensorLab(flag=0)])
)
test_salobj_dataloader = DataLoader(test_salobj_dataset,
batch_size=1,
shuffle=False,
num_workers=1)
# --------- 3. model define ---------
if(model_name=='u2net'):
print("...load U2NET---173.6 MB")
net = U2NET(3,1)
elif(model_name=='u2netp'):
print("...load U2NEP---4.7 MB")
net = U2NETP(3,1)
if torch.cuda.is_available():
net.load_state_dict(torch.load(model_dir))
net.cuda()
else:
net.load_state_dict(torch.load(model_dir, map_location='cpu'))
net.eval()
# --------- 4. inference for each image ---------
for i_test, data_test in enumerate(test_salobj_dataloader):
# print("inferencing:",img_name_list[i_test].split(os.sep)[-1])
inputs_test = data_test['image']
inputs_test = inputs_test.type(torch.FloatTensor)
if torch.cuda.is_available():
inputs_test = Variable(inputs_test.cuda())
else:
inputs_test = Variable(inputs_test)
d1,d2,d3,d4,d5,d6,d7= net(inputs_test)
# normalization
pred = d1[:,0,:,:]
pred = normPRED(pred)
# save results to test_results folder
if not os.path.exists(prediction_dir):
os.makedirs(prediction_dir, exist_ok=True)
save_output(img_name_list[i_test],pred,prediction_dir, show=True)
# sys.exit(0)
del d1,d2,d3,d4,d5,d6,d7
六、附件
见附件
想理解更多的 AI 技术干货，欢送上华为云的 AI 专区，目前有 AI 编程 Python 等六大实战营供大家收费学习。（六大实战营 link:http://su.modelarts.club/qQB9）
附件.zip71.28KB
点击关注，第一工夫理解华为云陈腐技术~
正文完
