用pytorch搭建簡單的語義分割(可訓練自己的數據集)

本文轉載自查看原文 2021-02-08 11:28 673 圖像分割

置頂 BlackBack_ 2020-07-07 22:59:17

1377

文章標簽：深度學習 pytorch 神經網絡

版權

用pytorch搭建簡單的語義分割(可訓練自己的數據集)

靈感來源：https://blog.csdn.net/weixin_44791964/article/details/102979289

本博客的搭建的網絡源於這位博主采用的keras框架，不過基於本人電腦配置做了一些網絡層數的改動。部分代碼引用大佬的代碼，其余均為本人原創。

整體文件目錄下排放：

在這里插入圖片描述

1、編碼器Mobilenet：

這里也有大佬關於Mobilenet的博客Mobilenet的介紹。簡單來說Mobilenet利用深度卷積使得數據量大大減少，有助於配置較低的機器，也可以應用到手機上。

import torch from torch.nn import * from torch.nn.functional import relu6 #第一個卷積塊 class Conv_block(Module): def __init__(self,inplanes,outplanes,strides): super(Conv_block, self).__init__() self.zeropad=ZeroPad2d(padding=1) self.conv=Conv2d(inplanes,outplanes,kernel_size=3,stride=strides,padding=0) self.BN=BatchNorm2d(outplanes,momentum=0.1) # self.relu=ReLU() def forward(self,x): x=self.zeropad(x) x=self.conv(x) x=self.BN(x) # x=self.relu(x) x=relu6(x) return x #除了第一個卷積塊的后面的深度卷積塊 class depthwise_block(Module): def __init__(self,inplanes,outplanes,strides): super(depthwise_block, self).__init__() self.zeropad=ZeroPad2d(padding=1) self.DW=Conv2d(inplanes,inplanes, #深度卷積,輸入和輸出通道一致 kernel_size=3,stride=strides, padding=0,groups=inplanes, #groups=inplanes是實現深度卷積的重點 bias=False) self.BN_1=BatchNorm2d(inplanes,momentum=0.1) self.BN_2=BatchNorm2d(outplanes,momentum=0.1) self.conv=Conv2d(inplanes,outplanes,kernel_size=1,stride=1) # self.relu=ReLU() def forward(self,x): x=self.zeropad(x) x=self.DW(x) x=self.BN_1(x) # x=self.relu(x) x = relu6(x) x=self.conv(x) x=self.BN_2(x) # x=self.relu(x) x=relu6(x) return x class Mobilenet(Module): cfg_filter=[32,64,128,128,256,256] #每個block的inplanes、outplanes cfg_stride=[1,2,1,2,1] #每個block的strides cfg_block=[] #初始化后的block集成一個列表 layer_data=[] #每個block處理后的output def __init__(self): super(Mobilenet, self).__init__() self.conv_block=Conv_block(3,32,2) #第一個conv block self.block_1=depthwise_block(32,64,1) self.block_2=depthwise_block(64,128,2) self.block_3=depthwise_block(128,128,1) self.block_4=depthwise_block(128,256,2) self.block_5=depthwise_block(256,256,1) def forward(self,inputs): x=inputs x=self.conv_block(x) x=self.block_1(x) x=self.block_2(x) x=self.block_3(x) x=self.block_4(x) x=self.block_5(x) return x #測試encoder網絡 if __name__ =="__main__": model=Mobilenet() inputs=torch.randn(1,416,416,3).permute(0,3,1,2) # inputs=torch.randn(1,3,416,416) # layers_list=model(inputs) outputs = model(inputs) print("layers_3 shape:" ) # print(layers_list[2].shape) print(outputs.shape)

2、解碼器Segnet：

解碼器對應着上面的編碼器，目的是把獲得的特征重新映射到比較搭的圖片中的每一個像素點，用於每一個像素點的分類。放大倍數上，和大佬博客不一樣的是，本人把最終的size放大到放入網絡的size，個人認為這樣有助於每個像素的特征得到對應。

import torch import numpy as np from torch.nn import * from torch.nn import functional as F from mobilenet_ import Mobilenet class Segnet(Module): cfg_filter=[256,128,64,32] conv_block=[] BN_block=[] def __init__(self,num_classes): super(Segnet, self).__init__() self.zeropad=ZeroPad2d(padding=1) self.conv_1=Conv2d(256,256,kernel_size=3,padding=0) self.conv_2=Conv2d(32,num_classes,kernel_size=3,padding=1) self.BN_1=BatchNorm2d(256,momentum=0.1) self.upsample=Upsample(scale_factor=2) for i in range(len(self.cfg_filter)-1): self.conv_block += [Conv2d(self.cfg_filter[i], self.cfg_filter[i + 1], kernel_size=3, padding=0)] self.BN_block +=[BatchNorm2d(self.cfg_filter[i+1])] self.conv_block=ModuleList(self.conv_block) self.BN_block = ModuleList(self.BN_block) def forward(self,o): #input:52,52,256 o=self.zeropad(o) o=self.conv_1(o) o=self.BN_1(o) #input:104,104,256 for j in range(3): o=self.upsample(o) o=self.zeropad(o) o=self.conv_block[j](o) o=self.BN_block[j](o) outputs=self.conv_2(o) return outputs #編碼器和解碼器組合 class Airplanesnet(Module): def __init__(self,classes1,BATCH_SIZE): super(Airplanesnet, self).__init__() self.encoder_part=Mobilenet() #Mobilenet()是從另一個py文件import過來的類 self.decoder_part=Segnet(classes1) self.classes=classes1 self.batch_size=BATCH_SIZE def forward(self,input_1): x=self.encoder_part(input_1) x=self.decoder_part(x) return x #測試decoder網絡 if __name__ =="__main__": model=Airplanesnet(classes1=2,BATCH_SIZE=1) inputs_1=torch.Tensor(torch.randn(1,3,416,416)) outputs_1=model(inputs_1) # outputs=outputs[3] print("outputs shape:" ) print(outputs_1.shape)

Segnet最后沒有進行softmax，因為training用的是CrossEntropyLoss。

3、訓練自己的數據集training：

本博客采用的是VOC2012語義分割兩類數據集，並通過自己的代碼處理以及數據增強得到看起來為黑色的，**每個像素值非0即1(目標像素值為1)**的label。

$[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-7j3otUTf-1594133488095)(C:\Users\86152\Desktop\博客\1.png)]$

$[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-WjVC1yuD-1594133488096)(C:\Users\86152\Desktop\博客\2.png)]$

代碼如下：

import torch import cv2 import os import argparse import numpy as np from PIL import Image from torch.nn import * from torch.optim import Adam from torch.utils.data import Dataset,DataLoader from segnet_ import Airplanesnet BATCH_SIZE1=1 #訓練的batch_size BATCH_SIZE2=1 #驗證的batch_size NUM_CLASSES=2 #分割的種類數 LR=1e-4 #學習率 EPOCH=20 #迭代次數 parser = argparse.ArgumentParser() parser.add_argument('--gpu',action='store_true',default=True,help='whether use gpu') parser.add_argument('--train_txt', type=str, default='D:/untitled/.idea/SS_torch/dataset/train.txt', help='about trian') parser.add_argument('--val_txt', type=str, default='D:/untitled/.idea/SS_torch/dataset/val.txt', help='about validation') opt = parser.parse_args() print(opt) txt_1 = opt.train_txt txt_2 = opt.val_txt #自定義數據集的類 class AirplanesDataset(Dataset): def __init__(self,txt_path): super(AirplanesDataset, self).__init__() paths=open("%s" % txt_path,"r") data=[] for lines in paths: path=lines.rstrip('\n') data.append(path) self.data=data self.len=len(data) def __getitem__(self, index): image=cv2.imread("D:/untitled/.idea/SS_torch/dataset/jpg_right/%s" %self.data[index]+".jpg",-1) label = cv2.imread("D:/untitled/.idea/SS_torch/dataset/png_right/%s"%self.data[index] +".png" , -1) image = cv2.resize(image, dsize=(416, 416)) label = cv2.resize(label, dsize=(416, 416)) image=torch.from_numpy(image) label=torch.from_numpy(label) image = image / 255.0 #歸一化 label[label>=0.5]=1 #label被resize后像素值會改變,調整像素值為原來的兩類 label[label < 0.5] = 0 image=image.permute(2,0,1) #調整圖像維度,方便載入model return image,label def __len__(self): return self.len train_dataset = AirplanesDataset(txt_1) # 訓練集 # 加載訓練數據集,並且分好mini-batch train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE1, shuffle=True) criterion = CrossEntropyLoss() # Loss model=Airplanesnet(NUM_CLASSES,BATCH_SIZE1) optimizer = Adam(model.parameters(), # 優化器 lr=LR) device=torch.device("cuda:0"if torch.cuda.is_available() else "cpu") #檢測是否有GPU加速 model.to(device) #網絡放入GPU里加速 model.load_state_dict(torch.load('D:/untitled/.idea/SS_torch/weights/SS_weight_2.pth')) #train函數 def train(epoch): running_loss=0.0 for batch_idx,data in enumerate(train_loader,0): #0是表示從0開始 image,label=data image,label=image.to(device),label.to(device) #數據放進GPU里 optimizer.zero_grad() #優化器參數清零 #forword+backward+update image=image.type(torch.FloatTensor) #轉化數據類型,不轉則會報錯 image=image.to(device) outputs=model(image) loss=criterion(outputs,label.long()) #進行loss計算 lll=label.long().cpu().numpy() #把label從GPU放進CPU loss.backward(retain_graph=True) #反向傳播(求導) optimizer.step() #優化器更新model權重 running_loss+=loss.item() #收集loss的值 if batch_idx % 100 ==99: print('[epoch: %d,idex: %2d] loss:%.3f' % (epoch+1,batch_idx+1,running_loss/322)) #訓練集的數量,可根據數據集調整 runing_loss=0.0 #收集的loss值清零 torch.save(model.state_dict(),f='D:/untitled/.idea/SS_torch/weights/SS_weight_3.pth') #保存權重 for epoch in range(EPOCH): #迭代次數 train(epoch)

4、預測文件predict：

把待測圖像放入samples文件夾中，輸出結果在outputs文件夾中

$[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-KxaWFVZU-1594133488097)(C:\Users\86152\Desktop\博客\4.png)]$

$[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-ucszxAAb-1594133488098)(C:\Users\86152\Desktop\博客\5.png)]$

predict的重點是：圖像從model輸出后得到的predict圖像經過【第70行pr=predict.argmax(axis=-1)】壓縮成一層，每個像素值為種類概率最高的該層的index，再遍歷全部像素點與種類index進行匹配，匹配成功則塗上上對應的顏色。

from segnet_ import Airplanesnet from PIL import Image import numpy as np import torch import argparse import cv2 import copy import os parser = argparse.ArgumentParser() parser.add_argument('--samples', type=str, default='D:/untitled/.idea/SS_torch/samples', help='samples') parser.add_argument('--outputs', type=str, default='D:/untitled/.idea/SS_torch/outputs', help='outputs') parser.add_argument('--weights', type=str, default='D:/untitled/.idea/SS_torch/weights/SS_weight_3.pth', help='weights') opt = parser.parse_args() print(opt) colors = [[0,0,0],[255,0,0]] NCLASSES = 2 BATCH_SIZE=1 img_way=opt.samples img_save=opt.outputs device=torch.device("cuda:0"if torch.cuda.is_available() else "cpu") #檢測是否有GPU加速 model=Airplanesnet(NCLASSES,BATCH_SIZE) #初始化model model.load_state_dict(torch.load(opt.weights)) #加載權重 model.to(device) #放入GPU for jpg in os.listdir(r"%s" %img_way): name = jpg[:-4] with torch.no_grad(): image=cv2.imread("%s" % img_way + "/" + jpg) old_image = copy.deepcopy(image) old_image = np.array(old_image) orininal_h = image.shape[0] #讀取的圖像的高 orininal_w = image.shape[1] #讀取的圖像的寬 方便之后還原大小 image = cv2.resize(image, dsize=(416, 416)) #調整大小 image = image / 255.0 #圖像歸一化 image = torch.from_numpy(image) image = image.permute(2, 0, 1) #顯式的調轉維度 image = torch.unsqueeze(image, dim=0) #改變維度,使得符合model input size image = image.type(torch.FloatTensor) #數據轉換,否則報錯 image = image.to(device) #放入GPU中計算 predict = model(image).cpu() # print(predict.shape) predict = torch.squeeze(predict) #[1,1,416,416]---->[1,416,416] predict =predict.permute(1, 2, 0) # print(jpg) predict = predict.numpy() # print(predict.shape) pr=predict.argmax(axis=-1) #把class數量的層壓縮為一層,Z軸上的值概率最高的返回該層index seg_img = np.zeros((416, 416,3)) #創造三層0矩陣,方便進行塗色匹配 #進行染色 for c in range(NCLASSES): seg_img[:, :, 0] += ((pr[:, :] == c) * (colors[c][0])).astype('uint8') seg_img[:, :, 1] += ((pr[:, :] == c) * (colors[c][1])).astype('uint8') seg_img[:, :, 2] += ((pr[:, :] == c) * (colors[c][2])).astype('uint8') seg_img = cv2.resize(seg_img,(orininal_w,orininal_h)) seg_img = np.array(seg_img) # 原圖和效果圖疊加 result = cv2.addWeighted(seg_img, 0.3, old_image, 0.7, 0., old_image, cv2.CV_32F) cv2.imwrite("%s/%s" % (img_save, name) + ".jpg", result) print("%s.jpg ------>done!!!" % name)

預測結果：

$[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-anQ4GWJE-1594133488099)(C:\Users\86152\Desktop\博客\6.jpg)]$

5、語義分割mIoU評測指標：

源碼：https://www.cnblogs.com/Trevo/p/11795503.html，即把兩個矩陣進行mIoU評測，所以我們之后要做的很簡單，就是跟predict相似，img經過model后輸出predict，然后與label進行匹配。

from segnet_ import Airplanesnet import numpy as np import torch import argparse import copy import cv2 NCLASSES = 2 BATCH_SIZE=1 #文件的加載路徑 parser = argparse.ArgumentParser() parser.add_argument('--val_txt', type=str, default='D:/untitled/.idea/SS_torch/dataset/val.txt', help='about validation') parser.add_argument('--weights', type=str, default='D:/untitled/.idea/SS_torch/weights/SS_weight_3.pth', help='weights') opt = parser.parse_args() print(opt) txt_path = opt.val_txt weight=opt.weights __all__ = ['SegmentationMetric'] class SegmentationMetric(object): #計算mIoU、accuracy的類 def __init__(self, numClass): self.numClass = numClass self.confusionMatrix = np.zeros((self.numClass,) * 2) def pixelAccuracy(self): # return all class overall pixel accuracy # acc = (TP + TN) / (TP + TN + FP + TN) acc = np.diag(self.confusionMatrix).sum() / self.confusionMatrix.sum() acc = round(acc,5) return acc def classPixelAccuracy(self): # return each category pixel accuracy(A more accurate way to call it precision) # acc = (TP) / TP + FP classAcc = np.diag(self.confusionMatrix) / self.confusionMatrix.sum(axis=1) return classAcc def meanPixelAccuracy(self): classAcc = self.classPixelAccuracy() meanAcc = np.nanmean(classAcc) return meanAcc def meanIntersectionOverUnion(self): # Intersection = TP Union = TP + FP + FN # IoU = TP / (TP + FP + FN) intersection = np.diag(self.confusionMatrix) union = np.sum(self.confusionMatrix, axis=1) + np.sum(self.confusionMatrix, axis=0) - np.diag( self.confusionMatrix) IoU = intersection / union mIoU = np.nanmean(IoU) mIoU =round(mIoU,4) return mIoU def genConfusionMatrix(self, imgPredict, imgLabel): # remove classes from unlabeled pixels in gt image and predict mask = (imgLabel >= 0) & (imgLabel < self.numClass) label = self.numClass * imgLabel[mask] + imgPredict[mask] count = np.bincount(label, minlength=self.numClass ** 2) confusionMatrix = count.reshape(self.numClass, self.numClass) return confusionMatrix def Frequency_Weighted_Intersection_over_Union(self): # FWIOU = [(TP+FN)/(TP+FP+TN+FN)] *[TP / (TP + FP + FN)] freq = np.sum(self.confusion_matrix, axis=1) / np.sum(self.confusion_matrix) iu = np.diag(self.confusion_matrix) / ( np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) - np.diag(self.confusion_matrix)) FWIoU = (freq[freq > 0] * iu[freq > 0]).sum() return FWIoU def addBatch(self, imgPredict, imgLabel): assert imgPredict.shape == imgLabel.shape self.confusionMatrix += self.genConfusionMatrix(imgPredict, imgLabel) def reset(self): self.confusionMatrix = np.zeros((self.numClass, self.numClass)) #讀取val.txt中的圖片的名稱 paths = open("%s" % txt_path, "r") data = [] for lines in paths: path = lines.rstrip('\n') data.append(path) device=torch.device("cuda:0"if torch.cuda.is_available() else "cpu") #檢測是否有GPU加速 model=Airplanesnet(NCLASSES,BATCH_SIZE) #初始化model model.load_state_dict(torch.load(opt.weights)) #加載權重 model.to(device) sum_1 = 0 # 累加每張圖片val的accuracy sum_2 = 0 # 累積每張圖片Val的mIoU for i in range(len(data)): image = cv2.imread("D:/untitled/.idea/SS_torch/dataset/jpg_right/%s" % data[i] + ".jpg", -1) label = cv2.imread("D:/untitled/.idea/SS_torch/dataset/png_right/%s" % data[i] + ".png", -1) orininal_h = image.shape[0] # 讀取的圖像的高 orininal_w = image.shape[1] # 讀取的圖像的寬 image = cv2.resize(image, dsize=(416, 416)) label = cv2.resize(label, dsize=(416, 416)) label[label >= 0.5] = 1 #label被resize后像素值會改變,調整像素值為原來的兩類 label[label < 0.5] = 0 image = image / 255.0 # 圖像歸一化 image = torch.from_numpy(image) image = image.permute(2, 0, 1) # 顯式的調轉維度 image = torch.unsqueeze(image, dim=0) # 改變維度,使得符合model input size image = image.type(torch.FloatTensor) # 數據轉換,否則報錯 image = image.to(device) # 放入GPU中計算 predict = model(image).cpu() predict = torch.squeeze(predict) # [1,1,416,416]---->[1,416,416] predict = predict.permute(1, 2, 0) predict = predict.detach().numpy() prc = predict.argmax(axis=-1) #進行mIoU和accuracy的評測 imgPredict =prc imgLabel = label metric = SegmentationMetric(2) metric.addBatch(imgPredict, imgLabel) acc = metric.pixelAccuracy() sum_1+=acc mIoU = metric.meanIntersectionOverUnion() sum_2+=mIoU print("%s.jpg :" % data[i]) print("accuracy: "+str(acc*100)+" %") print("mIoU: " +str(mIoU)) print("-------------------") # 全部圖片平均的accuracy和mIoU sum_1=sum_1/len(data) sum_2=sum_2/len(data) sum_1 = round(sum_1,5) sum_2 = round(sum_2,4) print("M accuracy: "+str(sum_1*100)+" %") print("M mIoU: " +str(sum_2))

評測結果：

$[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-qSqlYycr-1594133488100)(C:\Users\86152\Desktop\博客\6.png)]$

6、數據增強：

本語義分割代碼采用python的Augmentor庫，但是有個缺點就是每次增強只能一張圖片，多於一張會讓label和image的形變不對應，所以代碼有點繞，即讀取一張圖片，增強后用os把圖片移除，再把處理好的label和image分別放入不同文件夾，方便以上一系列操作。

import os import cv2 import argparse import Augmentor #文件路徑 parser = argparse.ArgumentParser() parser.add_argument('--Images', type=str, default='D:/untitled/.idea/SS_torch/Augmentor_img', help='true picture') parser.add_argument('--final', type=str, default='D:/untitled/.idea/SS_torch/Augmentor_img/output', help='final picture') parser.add_argument('--Masks', type=str, default='D:/untitled/.idea/SS_torch/Augmentor_mask', help='Mask picture') parser.add_argument('--jpg_right', type=str, default='D:/untitled/.idea/SS_torch/dataset/jpg_right', help='final picture') parser.add_argument('--png_right', type=str, default='D:/untitled/.idea/SS_torch/dataset/png_right', help='final masks') parser.add_argument('--transtxt', type=str, default='D:/untitled/.idea/SS_torch/dataset/trans.txt', help='transtxt') opt = parser.parse_args() print(opt) txt=opt.transtxt paths = open("%s" % txt, "r") data = [] for lines in paths: path = lines.rstrip('\n') data.append(path) imgway_1=opt.Images imgway_2=opt.final JPG_RIGHT=opt.jpg_right PNG_RIGHT=opt.png_right #for循環命名需要 n1 = 1 n2 = 1 #進行數據增強 for index in range(len(data)): #讀取需要增強的image和label image = cv2.imread("D:/untitled/.idea/SS_torch/dataset/jpg/%s" % data[index] + ".jpg", -1) mask = cv2.imread("D:/untitled/.idea/SS_torch/dataset/png/%s" % data[index] + ".png", -1) #保存至數據增強指定的文件夾中 cv2.imwrite("%s/%s.jpg" % (imgway_1, data[index]) ,image) cv2.imwrite("%s/%s.jpg" % (opt.Masks, data[index]) , mask) #數據增強主體 p = Augmentor.Pipeline(opt.Images) #讀取image p.ground_truth(opt.Masks) #讀取label,使得label和對應的image進行相同變化的augmentor p.rotate(probability=1, max_left_rotation=5, max_right_rotation=5) #旋轉圖片,左邊最大旋轉度,右邊最大旋轉度 p.shear(probability=1,max_shear_left=15,max_shear_right=15) #隨機區域形變 p.flip_left_right(probability=0.5) #按概率左右翻轉 p.zoom_random(probability=0.5, percentage_area=0.8) #按概率放大圖片 p.flip_top_bottom(probability=0.5) #按概率上下翻轉 p.sample(3) #產生3張圖片 os.remove("%s/%s.jpg"%(imgway_1,data[index])) #去除原來的img,防止mask和img不匹配 os.remove("%s/%s.jpg" % (opt.Masks, data[index])) #去除原來的mask,防止mask和img不匹配 #將數據增強后的img和mask進行對應改名並移動到制定的文件夾中 for filename in os.listdir(r"%s" % imgway_2): name = filename[:9] if name =="Augmentor": #該圖片是image name_1 = [] # 把image的數字名稱放入列表 name_1.append(filename[23:34]) #截取數字+格式 img = cv2.imread("%s" % imgway_2 + "/" + filename,-1) name1_1 = name_1[0] name2_1 = name1_1[:-6]+str(n1)+ name1_1[6:] #圖片在原來名稱基礎上改名 cv2.imwrite("%s/%s" % (JPG_RIGHT, name2_1 )+".jpg", img) n1+=1 if n1==4: #防止改名出現錯誤 n1=1 else: #該圖片是mask name_2 = [] # 把mask的數字名稱放入列表 name_2.append(filename[31:42]) #截取數字+格式 img_2 = cv2.imread("%s" % imgway_2 + "/" + filename, -1) name1_2 = name_2[0] name2_2 = name1_2[:-6] + str(n2) + name1_2[6:] #圖片在原來名稱基礎上改名 cv2.imwrite("%s/%s" % (PNG_RIGHT, name2_2)+".png", img_2) n2 += 1 if n2==4: #防止改名出現錯誤 n2=1

7、關於訓練集和驗證集的txt文件：

$[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-7ANI9VMS-1594133488100)(C:\Users\86152\Desktop\博客\7.png)]$

$[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-bfYuc0xh-1594133488101)(C:\Users\86152\Desktop\博客\8.png)]$

import os import random val_percent = 0.1 train_percent = 0.9 imagepath = 'dataset/jpg_right' txtsavepath = 'dataset' total_img = os.listdir(imagepath) num = len(total_img) list=range(num) tv = int(num * val_percent) #驗證個數 tr = int(num-tv) #訓練個數 num_trainval = random.sample(list, tv) #隨機獲取tv個片段 num_train = random.sample(list, tr) #隨機獲取tr個片段 ftrain = open('dataset/train.txt', 'w') fval = open('dataset/val.txt', 'w') for i in range(num): name = total_img[i][:-4] + '\n' #提取名字+轉行 if i in num_train: ftrain.write(name) else: fval.write(name) print("True") print(i+1) ftrain.close() fval.close()

總結：

該算法不使用預訓練模型，原因有二：第一，難以找到mobilenet的預訓練模型，並且還要進行個性化的修改，較麻煩。第二，模型數據量小，訓練好的權重大小為4MB左右(圖片數為357)，從頭開始訓練的效果也達到及格線。

$[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-egevJtTR-1594133488101)(C:\Users\86152\Desktop\博客\9.png)]$

免責聲明！

本站轉載的文章為個人學習借鑒使用，本站對版權不負任何法律責任。如果侵犯了您的隱私權益，請聯系本站郵箱yoyou2525@163.com刪除。

猜您在找 語義分割與數據集 PyTorch中的MIT ADE20K數據集的語義分割 pytorch—定義自己的數據集及加載訓練 Pytorch_加載自己的數據集並訓練視頻語義分割基准數據集與評估方法 pytorch_基於cifar創建自己的數據集並訓練【pytorch-ssd目標檢測】訓練自己創建的數據集教你使用pytorch-yolo訓練KITTI數據集！ facebook-pytorch maskrcnn-benchmark訓練自己的數據集基於pytorch實現Resnet對本地數據集的訓練