目錄
摘要
制作資料集
匯入所需要的庫
設定全域引數
影像預處理
讀取資料
設定模型
設定訓練和驗證
測驗
摘要
給你一張動漫圖片，你能告訴我出自哪個動漫嗎？今天我們就用GoogLeNet做到這一點，我選擇的的動漫有秦嶺神樹、吞噬星空和秦時明月，圖片樣例如下：
從上面的樣例空可以看出，不同的動漫風格差異還是很大的，下面就說說我如何實作動漫分類的，
制作資料集
制作資料集，需要用到從動漫中抽取圖片，抽取圖片需要用到ffmpeg工具，具體的安裝和使用參考：Windows安裝ffmpeg，使用ffmpeg從視頻中的抽取影像_AI浩-CSDN博客，
然后我們使用python，呼叫ffmeg實作對動漫視頻的批量抽取的邏輯，在D盤中新建兩個檔案夾，一個檔案夾命名qlss，這個檔案夾存放動漫視頻，一個檔案夾是imags,用于存放抽取的圖片，Python的代碼如下：
import os
# 創建三個串列用來存盤視頻檔案以及視頻地址
file_list = []
file_list_path=[]
filelist = []
# 源檔案目錄
dir_path = 'D:\\qlss'
#cmd命令存入str字串
str = 'ffmpeg ' + '-i {} -ss 00:00:30 -f image2 -vf fps=1/5 -qscale:v 2 ../imags/img_{}%05d.jpg'
filenameList=os.listdir(dir_path)
print(filenameList)
for j in range(len(filenameList)):
filepath=os.path.join(dir_path,filenameList[j])
str_cmd = str.format(filepath, filenameList[j].split('.')[0])
print(str_cmd)
os.popen(str_cmd)
上面代碼的思路：遍歷視頻檔案，使用Python在cmd中執行ffmpeg命令實作對圖片抽取，
新建專案
新建一個影像分類的專案，data里面放資料集，dataset檔案夾中自定義資料的讀取方法，這次我不采用默認的讀取方式，太簡單沒啥意思，然后再新建train.py和test.py
在專案的根目錄新建train.py，然后在里面寫訓練代碼，
匯入所需要的庫
我這次選用的模型是inception_v3，
import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
from dataset.dataset import SeedlingData
from torch.autograd import Variable
from torchvision.models import inception_v3
設定全域引數
設定BatchSize、學習率和epochs，判斷是否有cuda環境，如果沒有設定為cpu，
# 設定全域引數
modellr = 1e-4
BATCH_SIZE = 32
EPOCHS = 10
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')影像預處理
在做影像與處理時，train資料集的transform和驗證集的transform分開做，train的影像處理出了resize和歸一化之外，還可以設定影像的增強，比如旋轉、隨機擦除等一系列的操作，驗證集則不需要做影像增強，另外不要盲目的做增強，不合理的增強手段很可能會帶來負作用，甚至出現Loss不收斂的情況，注：inception_v3模型輸入的size是3×299×299
transform = transforms.Compose([
transforms.Resize((299, 299)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
transform_test = transforms.Compose([
transforms.Resize((299, 299)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])讀取資料
將資料集解壓后放到data檔案夾下面，如圖：
然后我們在dataset檔案夾下面新建 __init__.py和dataset.py，在dataset.py檔案夾寫入下面的代碼：
說一下代碼的核心邏輯，
第一步 建立字典，定義類別對應的ID,用數字代替類別，
第二步 在__init__里面撰寫獲取圖片路徑的方法，測驗集只有一層路徑直接讀取，訓練集在train檔案夾下面是類別檔案夾，先獲取到類別，再獲取到具體的圖片路徑，然后使用sklearn中切分資料集的方法，按照7：3的比例切分訓練集和驗證集，
第三步 在__getitem__方法中定義讀取單個圖片和類別的方法，由于影像中有位深度32位的，所以我在讀取影像的時候做了轉換，
# coding:utf8
import os
from PIL import Image
from torch.utils import data
from torchvision import transforms as T
from sklearn.model_selection import train_test_split
Labels={'吞噬星空': 0, '秦嶺神樹': 1, '秦時明月': 2}
class SeedlingData (data.Dataset):
def __init__(self, root, transforms=None, train=True, test=False):
"""
主要目標： 獲取所有圖片的地址，并根據訓練，驗證，測驗劃分資料
"""
self.test = test
self.transforms = transforms
if self.test:
imgs = [os.path.join(root, img) for img in os.listdir(root)]
self.imgs = imgs
else:
imgs_labels = [os.path.join(root, img) for img in os.listdir(root)]
imgs = []
for imglable in imgs_labels:
for imgname in os.listdir(imglable):
imgpath = os.path.join(imglable, imgname)
imgs.append(imgpath)
trainval_files, val_files = train_test_split(imgs, test_size=0.3, random_state=42)
if train:
self.imgs = trainval_files
else:
self.imgs = val_files
def __getitem__(self, index):
"""
一次回傳一張圖片的資料
"""
img_path = self.imgs[index]
img_path=img_path.replace("\\",'/')
if self.test:
label = -1
else:
labelname = img_path.split('/')[-2]
label = Labels[labelname]
data = Image.open(img_path).convert('RGB')
data = self.transforms(data)
return data, label
def __len__(self):
return len(self.imgs)
然后我們在train.py呼叫SeedlingData讀取資料 ，記著匯入剛才寫的dataset.py(from dataset.dataset import SeedlingData)
dataset_train = SeedlingData('data/train', transforms=transform, train=True)
dataset_test = SeedlingData("data/train", transforms=transform_test, train=False)
# 讀取資料
print(dataset_train.imgs)
# 匯入資料
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)
設定模型
使用CrossEntropyLoss作為loss，模型采用inception_v3，選用預訓練模型，更改全連接層，將最后一層類別設定為3，然后將模型放到DEVICE，優化器選用Adam，
# 實體化模型并且移動到GPU
criterion = nn.CrossEntropyLoss()
model_ft = inception_v3(pretrained=True)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, 3)
model_ft.to(DEVICE)
# 選擇簡單暴力的Adam優化器，學習率調低
optimizer = optim.Adam(model_ft.parameters(), lr=modellr)
def adjust_learning_rate(optimizer, epoch):
"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
modellrnew = modellr * (0.1 ** (epoch // 50))
print("lr:", modellrnew)
for param_group in optimizer.param_groups:
param_group['lr'] = modellrnew
設定訓練和驗證
inception_v3模型較早，model輸出有兩個引數，在這里要注意，
output,hid = model(data)
# 定義訓練程序
def train(model, device, train_loader, optimizer, epoch):
model.train()
sum_loss = 0
total_num = len(train_loader.dataset)
print(total_num, len(train_loader))
for batch_idx, (data, target) in enumerate(train_loader):
data, target = Variable(data).to(device), Variable(target).to(device)
output,hid = model(data)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print_loss = loss.data.item()
sum_loss += print_loss
if (batch_idx + 1) % 10 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
100. * (batch_idx + 1) / len(train_loader), loss.item()))
ave_loss = sum_loss / len(train_loader)
print('epoch:{},loss:{}'.format(epoch, ave_loss))
# 驗證程序
def val(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
total_num = len(test_loader.dataset)
print(total_num, len(test_loader))
with torch.no_grad():
for data, target in test_loader:
data, target = Variable(data).to(device), Variable(target).to(device)
output = model(data)
loss = criterion(output, target)
_, pred = torch.max(output.data, 1)
correct += torch.sum(pred == target)
print_loss = loss.data.item()
test_loss += print_loss
correct = correct.data.item()
acc = correct / total_num
avgloss = test_loss / len(test_loader)
print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
avgloss, correct, len(test_loader.dataset), 100 * acc))
# 訓練
for epoch in range(1, EPOCHS + 1):
adjust_learning_rate(optimizer, epoch)
train(model_ft, DEVICE, train_loader, optimizer, epoch)
val(model_ft, DEVICE, test_loader)
torch.save(model_ft, 'model.pth')完成上面的代碼就可以訓練了，這三部動漫的差別較大，所以很快就得到了非常好的結果，
測驗
我介紹兩種常用的測驗方式，第一種是通用的，通過自己手動加載資料集然后做預測，具體操作如下：
測驗集存放的目錄如下圖：
第一步 定義類別，這個類別的順序和訓練時的類別順序對應，一定不要改變順序！！！！
第二步 定義transforms，transforms和驗證集的transforms一樣即可，別做資料增強，
第三步 加載model，并將模型放在DEVICE里，
第四步 讀取圖片并預測圖片的類別，在這里注意，讀取圖片用PIL庫的Image，不要用cv2，transforms不支持，
import torch.utils.data.distributed
import torchvision.transforms as transforms
from PIL import Image
from torch.autograd import Variable
import os
# classes = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed',
# 'Common wheat','Fat Hen', 'Loose Silky-bent',
# 'Maize','Scentless Mayweed','Shepherds Purse','Small-flowered Cranesbill','Sugar beet')
classes=('吞噬星空','秦嶺神樹','秦時明月')
transform_test = transforms.Compose([
transforms.Resize((299, 299)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.load("model.pth")
model.eval()
model.to(DEVICE)
path='data/test/'
testList=os.listdir(path)
for file in testList:
img=Image.open(path+file)
img=transform_test(img)
img.unsqueeze_(0)
img = Variable(img).to(DEVICE)
out=model(img)
# Predict
_, pred = torch.max(out.data, 1)
print('Image Name:{},predict:{}'.format(file,classes[pred.data.item()]))選擇用了一集秦嶺神樹做測驗集，運行結果：
第二種 使用自定義的Dataset讀取圖片
import torch.utils.data.distributed
import torchvision.transforms as transforms
from dataset.dataset import SeedlingData
from torch.autograd import Variable
# classes = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed',
# 'Common wheat','Fat Hen', 'Loose Silky-bent',
# 'Maize','Scentless Mayweed','Shepherds Purse','Small-flowered Cranesbill','Sugar beet')
classes=('吞噬星空','秦嶺神樹','秦時明月')
transform_test = transforms.Compose([
transforms.Resize((299, 299)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.load("model.pth")
model.eval()
model.to(DEVICE)
dataset_test =SeedlingData('data/test/', transform_test,test=True)
print(len(dataset_test))
# 對應檔案夾的label
for index in range(len(dataset_test)):
item = dataset_test[index]
img, label = item
img.unsqueeze_(0)
data = Variable(img).to(DEVICE)
output = model(data)
_, pred = torch.max(output.data, 1)
print('Image Name:{},predict:{}'.format(dataset_test.imgs[index], classes[pred.data.item()]))
index += 1
完整的代碼和資料集上傳到CSDN：GoogLeNet影像分類.rar-深度學習檔案類資源-CSDN下載
轉載請註明出處，本文鏈接：https://www.uj5u.com/houduan/287161.html
標籤：python