#项目目录结构
|---blog
|---DataSet
|---dataset.py
|---models
|---AlexNet.py
|---DenseNet.py
|---GooleNet.py
|---VGG.py
|---pth
|---train.py
以下是每个文件的具体内容:
# dataset.py
from PIL import Image
import torch
import numpy as np
import os
from torchvision import transforms
class MyDataSet(torch.utils.data.Dataset):
def __init__(self, transform=None): # 初始化一些需要传入的参数
super(MyDataSet, self).__init__()
images = []
# 返回对应的图片和标签
image_list = []
label_list = []
cwd = 'F:/openvinortest/TEST/' # 数据根目录
classes = {'cat', 'dog'} #文件夹名称
for index, name in enumerate(classes):
class_path = cwd + name + '/'
for img_name in os.listdir(class_path):
image_list.append(class_path + img_name)
label_list.append(index)
print("There are %d class" % (len(image_list)))
# 组合文件
temp = np.array([image_list, label_list])
temp = temp.transpose() # 转置
# 对应的打乱顺序
np.random.shuffle(temp)
image_list = list(temp[:, 0])
label_list = list(temp[:, 1])
for i in range(len(image_list)):
images.append((image_list[i], int(label_list[i])))
self.imgs = images
self.transform = transform
def __getitem__(self, index):
fn, label = self.imgs[index]
img = Image.open(fn).convert('RGB')
if self.transform is not None:
img = self.transform(img)
else:
self.transform = transforms.Compose([transforms.Resize((227, 227)), transforms.ToTensor()])
img = self.transform(img)
return img, label
def __len__(self):
return len(self.imgs)
# 根据自己定义的那个勒MyDataset来创建数据集！注意是数据集！而不是loader迭代器
if __name__ == "__main__":
train_data = MyDataSet()
train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=64, shuffle=True)
for batch_index, data in train_loader:
print(len(batch_index))
# AlexNet.py
import torch.nn as nn
# 图片大小227
class AlexNet(nn.Module):
def __init__(self):
super(AlexNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11, stride=4), # 55*55 还是56
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # 27*27
nn.Conv2d(64, 192, kernel_size=5, padding=2), # 26 26
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # 13 13
nn.Conv2d(192, 384, kernel_size=3, padding=1), # 13 13
nn.ReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3, padding=1), # 13 13
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1), # 10 10
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # 5 5
)
self.classifier = nn.Sequential(
nn.Linear(6 * 6 * 256, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 2),
)
def forward(self, x):
x = self.features(x)
x = x.view(-1, 6 * 6 * 256)
x = self.classifier(x)
return x
# DenseNet.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict
class _DenseLayer(nn.Sequential):
def __init__(self, num_input_features, growth_rate, bn_size, drop_rate):
super(_DenseLayer, self).__init__()
self.add_module('norm1', nn.BatchNorm2d(num_input_features)),
self.add_module('relu1', nn.ReLU(inplace=True)),
self.add_module('conv1', nn.Conv2d(num_input_features, bn_size *
growth_rate, kernel_size=1, stride=1, bias=False)),
self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)),
self.add_module('relu2', nn.ReLU(inplace=True)),
self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate,
kernel_size=3, stride=1, padding=1, bias=False)),
self.drop_rate = drop_rate
def forward(self, x):
new_features = super(_DenseLayer, self).forward(x)
if self.drop_rate > 0:
new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)
return torch.cat([x, new_features], 1)
class _DenseBlock(nn.Sequential):
def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate):
super(_DenseBlock, self).__init__()
for i in range(num_layers):
layer = _DenseLayer(num_input_features + i * growth_rate, growth_rate, bn_size, drop_rate)
self.add_module('denselayer%d' % (i + 1), layer)
class _Transition(nn.Sequential):
def __init__(self, num_input_features, num_output_features):
super(_Transition, self).__init__()
self.add_module('norm', nn.BatchNorm2d(num_input_features))
self.add_module('relu', nn.ReLU(inplace=True))
self.add_module('conv', nn.Conv2d(num_input_features, num_output_features, kernel_size=1, stride=1, bias=False))
self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))
class DenseNet(nn.Module):
def __init__(self, growth_rate=32, block_config=(6, 12, 24, 16), num_init_features=64, bn_size=4, drop_rate=0,
num_classes=2):
super(DenseNet, self).__init__()
# First convolution
self.features = nn.Sequential(OrderedDict([
('conv0', nn.Conv2d(3, num_init_features, kernel_size=7, stride=2, padding=3, bias=False)),
('norm0', nn.BatchNorm2d(num_init_features)),
('relu0', nn.ReLU(inplace=True)),
('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)),
]))
# Each denseblock
num_features = num_init_features
for i, num_layers in enumerate(block_config):
block = _DenseBlock(num_layers=num_layers, num_input_features=num_features,
bn_size=bn_size, growth_rate=growth_rate, drop_rate=drop_rate)
self.features.add_module('denseblock%d' % (i + 1), block)
num_features = num_features + num_layers * growth_rate
if i != len(block_config) - 1:
trans = _Transition(num_input_features=num_features, num_output_features=num_features // 2)
self.features.add_module('transition%d' % (i + 1), trans)
num_features = num_features // 2
# Final batch norm
self.features.add_module('norm5', nn.BatchNorm2d(num_features))
# Linear layer
self.classifier = nn.Linear(num_features, num_classes)
# Official init from torch repo.
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal(m.weight.data)
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.bias.data.zero_()
def forward(self, x):
features = self.features(x)
out = F.relu(features, inplace=True)
out = F.avg_pool2d(out, kernel_size=7, stride=1).view(features.size(0), -1)
out = self.classifier(out)
return out
def densenet121(**kwargs):
model = DenseNet(num_init_features=64, growth_rate=32, block_config=(6, 12, 24, 16), **kwargs)
return model
def densenet169(**kwargs):
model = DenseNet(num_init_features=64, growth_rate=32, block_config=(6, 12, 32, 32), **kwargs)
return model
def densenet201(**kwargs):
model = DenseNet(num_init_features=64, growth_rate=32, block_config=(6, 12, 48, 32), **kwargs)
return model
def densenet161(**kwargs):
model = DenseNet(num_init_features=96, growth_rate=48, block_config=(6, 12, 36, 24), **kwargs)
return model
if __name__ == '__main__':
# 'DenseNet', 'densenet121', 'densenet169', 'densenet201', 'densenet161'
# Example
net = DenseNet()
print(net)
# GooleNet.py
import torch
import torch.nn as nn
import torch.nn.functional as F
# 编写卷积+bn+relu模块
class BasicConv2d(nn.Module):
def __init__(self, in_channels, out_channals, **kwargs):
super(BasicConv2d, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channals, **kwargs)
self.bn = nn.BatchNorm2d(out_channals)
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
return F.relu(x)
# 编写Inception模块
class Inception(nn.Module):
def __init__(self, in_planes,
n1x1, n3x3red, n3x3, n5x5red, n5x5, pool_planes):
super(Inception, self).__init__()
# 1x1 conv branch
self.b1 = BasicConv2d(in_planes, n1x1, kernel_size=1)
# 1x1 conv -> 3x3 conv branch
self.b2_1x1_a = BasicConv2d(in_planes, n3x3red, kernel_size=1)
self.b2_3x3_b = BasicConv2d(n3x3red, n3x3, kernel_size=3, padding=1)
# 1x1 conv -> 3x3 conv -> 3x3 conv branch
self.b3_1x1_a = BasicConv2d(in_planes, n5x5red, kernel_size=1)
self.b3_3x3_b = BasicConv2d(n5x5red, n5x5, kernel_size=3, padding=1)
self.b3_3x3_c = BasicConv2d(n5x5, n5x5, kernel_size=3, padding=1)
# 3x3 pool -> 1x1 conv branch
self.b4_pool = nn.MaxPool2d(3, stride=1, padding=1)
self.b4_1x1 = BasicConv2d(in_planes, pool_planes, kernel_size=1)
def forward(self, x):
y1 = self.b1(x)
y2 = self.b2_3x3_b(self.b2_1x1_a(x))
print(y2.shape)
y3 = self.b3_3x3_c(self.b3_3x3_b(self.b3_1x1_a(x)))
y4 = self.b4_1x1(self.b4_pool(x))
# y的维度为[batch_size, out_channels, C_out,L_out]
# 合并不同卷积下的特征图
return torch.cat([y1, y2, y3, y4], 1)
class GoogLeNet(nn.Module):
def __init__(self):
super(GoogLeNet, self).__init__()
self.pre_layers = BasicConv2d(3, 192, kernel_size=3, padding=1)
self.a3 = Inception(192, 64, 96, 128, 16, 32, 32)
self.b3 = Inception(256, 128, 128, 192, 32, 96, 64)
self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
self.a4 = Inception(480, 192, 96, 208, 16, 48, 64)
self.b4 = Inception(512, 160, 112, 224, 24, 64, 64)
self.c4 = Inception(512, 128, 128, 256, 24, 64, 64)
self.d4 = Inception(512, 112, 144, 288, 32, 64, 64)
self.e4 = Inception(528, 256, 160, 320, 32, 128, 128)
self.a5 = Inception(832, 256, 160, 320, 32, 128, 128)
self.b5 = Inception(832, 384, 192, 384, 48, 128, 128)
self.avgpool = nn.AvgPool2d(8, stride=1)
self.drop = nn.Dropout(0.4)
self.linear = nn.Linear(1024, 2)
def forward(self, x):
out = self.pre_layers(x)
out = self.a3(out)
out = self.b3(out)
out = self.maxpool(out)
out = self.a4(out)
out = self.b4(out)
out = self.c4(out)
out = self.d4(out)
out = self.e4(out)
out = self.maxpool(out)
out = self.a5(out)
out = self.b5(out)
out = self.avgpool(out)
out = self.drop(out)
out = out.view(out.size(0), -1)
out = self.linear(out)
return out
# VGG.py
import torch.nn as nn
# 需继承torch.nn.Module类
class VGG16(nn.Module):
def __init__(self):
super(VGG16, self).__init__()
# 定义卷积层和池化层，共13层卷积，5层池化
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
# 简化版全连接层
# self.classifier = nn.Sequential(
# nn.Linear(4 * 4 * 512, 1024),
# nn.ReLU(),
# nn.Dropout(p=0.5),
# nn.Linear(1024, 1024),
# nn.ReLU(),
# nn.Dropout(p=0.5),
# nn.Linear(1024, 2)
# )
# VGG-16的全连接层
self.classes = nn.Sequential(
nn.Linear(7 * 7 * 512, 4096),
nn.ReLU(),
nn.Dropout(p=0.5),
nn.Linear(4096, 4096),
nn.ReLU(),
nn.Dropout(p=0.5),
nn.Linear(4096, 2)
)
# 定义每次执行的计算步骤
def forward(self, x):
x = self.features(x)
x = x.view(-1, 7 * 7 * 512)
x = self.classes(x)
return x
# train.py
import torch
import time
import cv2
import os
import random
import math
import torchvision
import torch.nn.functional as F
import torch.nn as nn
import matplotlib.pyplot as plt
from DataSet.dataset import MyDataSet
from models.DenseNet import DenseNet
from models.AlexNet import AlexNet
epochNum = 100
batchSize = 16
# 模型类实例
model = AlexNet()
# 如果GPUs可用，则将模型上需要计算的所有参数复制到GPUs上
if torch.cuda.is_available():
model = model.cuda()
print('支持cuda,GPU')
# loss值
loss_fn = nn.CrossEntropyLoss()
# 模型优化器
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
def train():
for epoch in range(1, epochNum + 1):
train_data = MyDataSet()
train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=batchSize, shuffle=True,
drop_last=True)
# 损失值
running_loss = 0.0
# 预测的正确数
running_correct = 0
for batch, data in enumerate(train_loader, 1):
x, y = data
if torch.cuda.is_available():
x, y = x.cuda(), y.cuda()
outputs = model(x)
_, y_pred = torch.max(outputs.detach(), 1)
# 将Varibale的梯度置零
optimizer.zero_grad()
# 计算损失值
loss = loss_fn(outputs, y)
# 反向传播求导
loss.backward()
# 更新所有参数
optimizer.step()
running_loss += loss.detach().item()
running_correct += torch.sum(y_pred == y)
if batch % 10 == 0:
print('Batch {}/{},Train Loss:{:.2f},Train Acc:{:.2f}%'.format(
batch, 25000 / batchSize, running_loss / batch,
100 * running_correct.item() / (batchSize * batch)
))
epoch_loss = running_loss * batchSize / 25000
epoch_acc = 100 * running_correct.item() / 25000
print(' Loss:{:.2f} Acc:{:.2f}%'.format(epoch_loss, epoch_acc))
torch.save(model.state_dict(), './pth/' + str(batch) + '.pkl')
# train()
# 测试
def test():
dst = MyDataSet()
data_loader_test = torch.utils.data.DataLoader(dst, batch_size=1, shuffle=False)
model.eval()
# 加载模型
model.load_state_dict(torch.load('./pth/2.pkl'))
# 保存测试结果
results = []
count = 0
# tqdm模块用于显示进度条
for imgs, path in data_loader_test:
if torch.cuda.is_available():
X = imgs.cuda()
outputs = model(X)
# pred表示是哪个对象，0=cat，1=dog
# probability表示是否个对象的概率
probability, pred = torch.max(F.softmax(outputs, dim=1).detach(), dim=1)
if pred == path.cuda():
print("OK")
count = count + 1
print(count)
# test()
if __name__ == "__main__":
train()
test()