ModelArts Notebook快速開(kāi)源項(xiàng)目實(shí)戰(zhàn) — U2Net

ModelArts Notebook快速開(kāi)源項(xiàng)目實(shí)戰(zhàn) — U2Net

一、U2Net介紹

U2Net是一個(gè)優(yōu)秀的顯著性目標(biāo)檢測(cè)算法，由Qin Xuebin等人發(fā)表在Pattern Recognition 2020期刊[Arxiv]。U2Net名稱(chēng)的來(lái)源在于其網(wǎng)絡(luò)結(jié)構(gòu)由兩層嵌套的Unet結(jié)構(gòu)，可以在不需要預(yù)訓(xùn)練骨干網(wǎng)絡(luò)的情況下從零開(kāi)始訓(xùn)練，擁有優(yōu)異的表現(xiàn)。其網(wǎng)絡(luò)結(jié)構(gòu)如圖1所示。

項(xiàng)目開(kāi)源地址：https://github.com/xuebinqin/U-2-Net

二、創(chuàng)建Notebook開(kāi)發(fā)環(huán)境

進(jìn)入ModelArts控制臺(tái)

選擇開(kāi)發(fā)環(huán)境 -> Notebook -> 創(chuàng)建

創(chuàng)建Notebook

3.1 可以選擇和任務(wù)相關(guān)的名稱(chēng)，方便管理；

3.2 為了減少不必要的資源消耗，建議開(kāi)啟自動(dòng)停止；

3.3 U2Net所需的運(yùn)行環(huán)境在公共鏡像中已經(jīng)包含，可以選擇pytorch1.4-cuda10.1-cudnn7-ubuntu18.04；

3.4 建議選擇GPU類(lèi)型，方便模型快速訓(xùn)練；

3.5 選擇立即創(chuàng)建 -> 提交，等待notebook創(chuàng)建完成后打開(kāi)Notebook。

導(dǎo)入開(kāi)源項(xiàng)目源碼（git/手動(dòng)上傳）

4.1 在Terminal使用git克隆遠(yuǎn)程倉(cāng)庫(kù)

cd work # 注意：只有/home/ma-user/work目錄及其子目錄下的文件在Notebook實(shí)例關(guān)閉后會(huì)保存 git clone https://github.com/xuebinqin/U-2-Net.git

4.2 如果git速度較慢也可以從本地上傳代碼，直接將壓縮包拖到左側(cè)文件目錄欄或者采用OBS上傳。

三、 數(shù)據(jù)準(zhǔn)備

下載訓(xùn)練數(shù)據(jù)APDrawing dataset

使用Wget直接下載到Notebook，也可下載本地后再拖拽到Notebook中。

wget https://cg.cs.tsinghua.edu.cn/people/~Yongjin/APDrawingDB.zip unzip APDrawingDB.zip

注：如果數(shù)據(jù)集較大（>5GB）需要下載到其它目錄（實(shí)例停止后會(huì)被刪除），建議存放在OBS中，需要的時(shí)候隨時(shí)拉取。

#從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')

切分訓(xùn)練數(shù)據(jù)

數(shù)據(jù)集中./APDrawingDB/data/train中包含了420張訓(xùn)練圖片，分辨率為512*1024，左側(cè)為輸入圖像，右側(cè)為對(duì)應(yīng)的ground truth。我們需要將大圖從中間切分為兩個(gè)子圖。

2.1 在Notebook開(kāi)發(fā)環(huán)境中新建一個(gè)Pytorch-1.4的jupyter Notebook文件，名稱(chēng)可以為split.ipynb，腳本將會(huì)在./APDrawingDB/data/train/split目錄下生成840張子圖，其中原始圖像以.jpg結(jié)尾，gt圖像以.png結(jié)尾，方便后續(xù)訓(xùn)練代碼讀取【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')

將切分好的數(shù)據(jù)按照如下層級(jí)結(jié)構(gòu)整理出訓(xùn)練和測(cè)試所需的datasets文件夾

datasets/

├── test (70張切分圖片，只包含原圖)

└── train (840張切分圖片，包含420張?jiān)瓐D及對(duì)應(yīng)的gt)

注：可以將切分好的數(shù)據(jù)集保存到OBS目錄中，減少./work的磁盤(pán)空間占用。

完整的U-2-Net項(xiàng)目結(jié)構(gòu)如下所示：

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

四、訓(xùn)練

官方提供的訓(xùn)練代碼中數(shù)據(jù)的路徑和我們的datasets有些區(qū)別，需要對(duì)訓(xùn)練腳本進(jìn)行一些修改，建議使用jupyter notebook方便排除錯(cuò)誤

新建一個(gè)Pytorch-1.4的jupyter Notebook文件，名稱(chēng)可以為train.ipynb

import moxing as mox # 如果需要從OBS拷貝切分好的訓(xùn)練數(shù)據(jù) #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

五、測(cè)試

新建一個(gè)Pytorch-1.4的jupyter Notebook文件，名稱(chēng)可以為test.ipynb

import moxing as mox # 拷貝數(shù)據(jù) 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): #此函數(shù)可以在notebook中展示模型的預(yù)測(cè)效果 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

六、附件

見(jiàn)附件

想了解更多的AI技術(shù)干貨，歡迎上華為云的AI專(zhuān)區(qū)，目前有AI編程Python等六大實(shí)戰(zhàn)營(yíng)供大家免費(fèi)學(xué)習(xí)。（六大實(shí)戰(zhàn)營(yíng)link:http://su.modelarts.club/qQB9）

附件： 附件.zip 71.28KB 下載次數(shù)：21次

AI平臺(tái) AI開(kāi)發(fā)平臺(tái)ModelArts Jupyter notebook 機(jī)器學(xué)習(xí) 機(jī)器視覺(jué)

版權(quán)聲明：本文內(nèi)容由網(wǎng)絡(luò)用戶(hù)投稿，版權(quán)歸原作者所有，本站不擁有其著作權(quán)，亦不承擔(dān)相應(yīng)法律責(zé)任。如果您發(fā)現(xiàn)本站中有涉嫌抄襲或描述失實(shí)的內(nèi)容，請(qǐng)聯(lián)系我們jiasou666@gmail.com 處理，核實(shí)后本網(wǎng)站將在24小時(shí)內(nèi)刪除侵權(quán)內(nèi)容。

版權(quán)聲明：本文內(nèi)容由網(wǎng)絡(luò)用戶(hù)投稿，版權(quán)歸原作者所有，本站不擁有其著作權(quán)，亦不承擔(dān)相應(yīng)法律責(zé)任。如果您發(fā)現(xiàn)本站中有涉嫌抄襲或描述失實(shí)的內(nèi)容，請(qǐng)聯(lián)系我們jiasou666@gmail.com 處理，核實(shí)后本網(wǎng)站將在24小時(shí)內(nèi)刪除侵權(quán)內(nèi)容。