|
1.类的定义
import torch.nn as nn
import torch
class BasicBlock(nn.Module):
expansion = 1#是否可以调用
def __init__(self, in_channel, out_channel, stride=1, downsample=None, **kwargs):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,
kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(out_channel)
self.relu = nn.ReLU()
self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,
kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(out_channel)
self.downsample = downsample
def forward(self, x):
identity = x
if self.downsample is not None:#残差结构实虚判断
identity = self.downsample(x)#虚函数
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out += identity
out = self.relu(out)
return out
class ResNet(nn.Module):
def __init__(self,
block,#残差结构
blocks_num,#残差结构数目
num_classes=1000,#训练集分类数目
include_top=True,#复杂可选
groups=1,
width_per_group=64):
super(ResNet, self).__init__()
self.include_top = include_top
self.in_channel = 64#卷积核个数
self.groups = groups#不用
self.width_per_group = width_per_group#不用
self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2,#输入通道数 卷积核个数
padding=3, bias=False)
self.bn1 = nn.BatchNorm2d(self.in_channel)#通道个数和卷核个数一样
self.relu = nn.ReLU(inplace=True)#激活
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)#池化
self.layer1 = self._make_layer(block, 64, blocks_num[0])
self.layer2 = self._make_layer(block, 128, blocks_num[1], stride=2)
self.layer3 = self._make_layer(block, 256, blocks_num[2], stride=2)
self.layer4 = self._make_layer(block, 512, blocks_num[3], stride=2)
if self.include_top:#默认为Ture
self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) # output size = (1, 1)#平均池化
self.fc = nn.Linear(512 * block.expansion, num_classes)#全连接层 512*1 分类数
for m in self.modules():#卷积层初始化！！！
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
def _make_layer(self, block, channel, block_num, stride=1):#通过上面定义好的残差结构造层 残差结构 卷积核个数 残差结构个数 步长默认为一
downsample = None#默认实
if stride != 1 or self.in_channel != channel * block.expansion:#用不到
downsample = nn.Sequential(
nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(channel * block.expansion))
layers = []
layers.append(block(self.in_channel,#第一层虚线残差结构
channel,
downsample=downsample,
stride=stride,
groups=self.groups,
width_per_group=self.width_per_group))
self.in_channel = channel * block.expansion
for _ in range(1, block_num):#后面全是实线残差结构
layers.append(block(self.in_channel,
channel,
groups=self.groups,
width_per_group=self.width_per_group))
return nn.Sequential(*layers)#*列表去括号
def forward(self, x):#正向传播
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
if self.include_top:
x = self.avgpool(x)
x = torch.flatten(x, 1)#压平 变成一维矩阵
x = self.fc(x)#全连接
return x
def resnet18(num_classes=1000, include_top=True):
# https://download.pytorch.org/models/resnet34-333f7ec4.pth
return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes, include_top=include_top)
2.训练分类模型
?
import os
import json
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets
from tqdm import tqdm
from sun import resnet18
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")#优先使用Gpu0如果有 没有则cpu
print("using {} device.".format(device))
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224),#随机裁剪224*224
transforms.RandomHorizontalFlip(),#随机翻转
transforms.ToTensor(),#转化成矩阵吧 ！！
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),#标准化处理
"val": transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}
#data_root = os.path.abspath(os.path.join(os.getcwd(), "../..")) #os.getcwd()获取当前文件目录 ..返回上一级目录../..返回上两级目录 abspath()返回绝对路径
#data_root=r"D:\artificial intelligence\Cat and dog recognition\project_data"
#image_path = os.path.join(data_root, "train1") # flower data set path
image_trpath_folder=r"D:\artificial intelligence\Cat and dog recognition\project_data\train"
#assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
train_dataset = datasets.ImageFolder(root=image_trpath_folder,#图片文件夹加载
transform=data_transform["train"])
train_num = len(train_dataset)
# {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
flower_list = train_dataset.class_to_idx#获取图片类型的字典
cla_dict = dict((val, key) for key, val in flower_list.items())#将字典反过来
# write dict into json file
json_str = json.dumps(cla_dict, indent=4)#编码成json
with open('class_indices.json', 'w') as json_file:#写入json文件
json_file.write(json_str)
batch_size = 16#批处理数量处理
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
print('Using {} dataloader workers every process'.format(nw))
train_loader = torch.utils.data.DataLoader(train_dataset,#加载图片
batch_size=batch_size, shuffle=True,#批处理数量
# 16
num_workers=nw)#！！！
image_vapath_folder=r"D:\artificial intelligence\Cat and dog recognition\project_data\test"#测试文件夹加载
validate_dataset = datasets.ImageFolder(root=image_vapath_folder,#预处理
transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,#图片加载
batch_size=batch_size, shuffle=False,#不洗牌
num_workers=nw)
print("using {} images for training, {} images for validation.".format(train_num,
val_num))
net = resnet18()#resnet网络对象
# load pretrain weights
# download url: https://download.pytorch.org/models/resnet34-333f7ec4.pth
model_weight_path = "./resnet18_pre.pth"
#assert os.path.exists(model_weight_path), "file {} does not exist.".format(model_weight_path)
#net.load_state_dict(torch.load(model_weight_path, map_location=device))
# for param in net.parameters():
# param.requires_grad = False
# change fc layer structure
in_channel = net.fc.in_features
net.fc = nn.Linear(in_channel, 2)#全连接层
net.to(device)
# 损失函数
loss_function = nn.CrossEntropyLoss()
#
#params = [p for p in net.parameters() if p.requires_grad]
optimizer = optim.Adam(net.parameters(), lr=0.0001)#学习rate优化器
epochs = 20#训练次数
best_acc = 0.0#准确rate初始化
save_path = './resNet18.pth'#权重保存
train_steps = len(train_loader)
for epoch in range(epochs):#
# train
net.train()#训练
running_loss = 0.0
train_bar = tqdm(train_loader)#添加训练进度条 返回迭代器
"""
enumerate()
names = ["Alice","Bob","Carl"]
for index,value in enumerate(names):
print(f'{index}: {value}')
"""
for step, data in enumerate(train_bar):
images, labels = data
optimizer.zero_grad()#清空之前的梯度信息
logits = net(images.to(device))#这参数
loss = loss_function(logits, labels.to(device))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
epochs,
loss)
net.eval()# 启用dropout方法
acc = 0.0 # accumulate accurate number / epoch
with torch.no_grad():
val_bar = tqdm(validate_loader)
for val_data in val_bar:
val_images, val_labels = val_data
outputs = net(val_images.to(device))
# loss = loss_function(outputs, test_labels)
predict_y = torch.max(outputs, dim=1)[1]
acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
val_bar.desc = "valid epoch[{}/{}]".format(epoch + 1,
epochs)
val_accurate = acc / val_num
print('[epoch %d] train_loss: %.3f val_accuracy: %.3f' %
(epoch + 1, running_loss / train_steps, val_accurate))
if val_accurate > best_acc:
best_acc = val_accurate
torch.save(net.state_dict(), save_path)
print('Finished Training')
if __name__ == '__main__':
main()
3.利用训练好的模型进行图片的识别和分类?
import torch
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import json
from sun import resnet18
data_transform=transforms.Compose([
transforms.Resize((224,224)),
transforms.ToTensor(),
transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
])
preimage_path=r"D:\artificial intelligence\Cat and dog recognition\project_data\test\test1\182.jpg"
image=Image.open(preimage_path)
plt.imshow(image)#显示图片格式
image=data_transform(image)
image=torch.unsqueeze(image,dim=0)#扩充维度batch 第0个维度一张图片3通道 224*224
try:
json_file=open('./class_indices.json','r')
class_indict=json.load(json_file)
except Exception as e:
print(e)
exit(-1)
model=resnet18(num_classes=2)
model_weight_path="./resNet18.pth"
model.load_state_dict(torch.load(model_weight_path))#载入模型
model.eval()#关闭dropout
with torch.no_grad():#默认进行反向传播 这个不进行
output=torch.squeeze(model(image))#模型载入图片计算并压缩
pre_list=torch.softmax(output,dim=0)#在列上进行概率计算
pre=torch.argmax(pre_list).numpy()#取最大
print(class_indict[str(pre)],pre_list[pre].item())
plt.show()
总结：本人第一次写博文，希望大家多多支持，有不懂得地方可以找我交流或者评论区留言。
|