vgg16是一个预训练的图片对象识别模型，它能区分一千种类型的物体，根据识别到的物体返回编号，编号对应的物体名称存放在vgg16_imagenet_class.json中
vgg16的使用
python
import torch
import torchvision.transforms as transforms
from torchvision.models import vgg16
from PIL import Image
import json
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 设置设备为 GPU（如果可用）或 CPU
torch.cuda.is_available() # 检查是否有可用的 GPU
# 假设你已经下载好了 vgg16 的 pth 文件，并保存在指定路径
local_weights_path = 'vgg16.pth' # 本地预训练权重文件的路径
# 加载 VGG16 模型（不加载预训练权重）
model = vgg16(pretrained=False)
# 加载本地的预训练权重
model.load_state_dict(torch.load(local_weights_path))
model.eval() # 设置模型为评估模式
# 定义图像预处理步骤
preprocess = transforms.Compose([
transforms.Resize(256), # 调整图像大小
transforms.CenterCrop(224), # 中心裁剪
transforms.ToTensor(), # 转换为张量
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 标准化
])
# 加载并预处理图像
img_path = 'data/presidential_doggy_door/valid/not_bo/137.jpg' # 替换为你的图像路径
img = Image.open(img_path)
img_tensor = preprocess(img)
img_tensor = img_tensor.unsqueeze(0) # 添加批次维度
# 进行预测
with torch.no_grad():
output = model(img_tensor)
# 解释预测结果
# 加载 ImageNet 类别标签
with open('data/vgg16_imagenet_class.json') as f:
labels = json.load(f)
# 获取预测的类别
_, predicted_idx = torch.max(output, 1)
print('predicted_idx:',predicted_idx.item())
# predicted_label = labels[str(predicted_idx.item())]
# print(f'预测结果: {predicted_label}')
vgg16迁移训练微调
在下面的微调中，我冻结原模型参数，我在一千层后，添加一个全新连接层，指定输出类别是1，即只区分是名叫“bo”这只狗，如果是其它动物或者其它狗就不会是1，最终这个模型就会变成只区分这是不是“bo”这只狗
python
import torch
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms.v2 as transforms
# import torchvision.io as tv_io
import glob
# import json
from PIL import Image
# import utils
from torchvision.models import vgg16
from torchvision.models import VGG16_Weights
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 设置设备为 GPU（如果可用）或 CPU
torch.cuda.is_available() # 检查是否有可用的 GPU
# 加载预训练的 VGG16 网络（基于 ImageNet 数据集）
weights = VGG16_Weights.DEFAULT # 获取 VGG16 的默认预训练权重
# 假设你已经下载好了 vgg16 的 pth 文件，并保存在指定路径
local_weights_path = 'vgg16.pth' # 本地预训练权重文件的路径
# 加载 VGG16 模型
vgg_model = vgg16() # 初始化 VGG16 模型
# 加载本地的预训练权重
vgg_model.load_state_dict(torch.load(local_weights_path)) # 加载本地预训练权重
vgg_model.to(device) # 将模型移动到指定设备
# 只训练新的层，所以把之前的层训练选项关掉，True 为修改前面的参数，False 为不修改前面的参数
######################################################
vgg_model.requires_grad_(False) # 冻结 VGG16 模型的参数,这里决定是否修改原模型的参数
######################################################
print("VGG16 解冻") # 打印冻结信息
# 定义输出类别数，因为VGG16 模型的最后一层全连接层的输出维度是 1000（对应 ImageNet 数据集的 1000 个类别）
N_CLASSES = 1
my_model = nn.Sequential( # 构建新的模型，包含 VGG16 和一个新的全连接层
vgg_model,
nn.Linear(1000, N_CLASSES) # 添加一个新的全连接层，输出类别数为1
)
# 打印模型的参数是否被冻结
# for idx, param in enumerate(my_model.parameters()):
# print(idx, param.requires_grad)
loss_function = nn.BCEWithLogitsLoss() # 定义损失函数为二元交叉熵损失
# 定义优化器为 Adam, 可选参数学习率为 0.000001，因为我们只训练最后一层，所以学习率可以设置的很小
#学习率很小的作用是，不会对原模型的参数造成太大的影响，尤其是原模型的参数已经很好了，我们只需要进行微调
######################################################
optimizer = Adam(my_model.parameters(), lr=.000002)
######################################################
my_model = my_model.to(device) # 将模型移动到指定设备
pre_trans = weights.transforms() # 获取预训练权重的预处理变换
# 定义数据集
DATA_LABELS = ["bo", "not_bo"] # 定义数据标签
class MyDataset(Dataset): # 自定义数据集类
def __init__(self, data_dir): # 初始化函数，接受数据目录
self.imgs = [] # 初始化图像列表
self.labels = [] # 初始化标签列表
for l_idx, label in enumerate(DATA_LABELS): # 遍历每个标签
data_paths = glob.glob(data_dir + label + '/*.jpg', recursive=True) # 获取所有图像路径
for path in data_paths: # 遍历每个图像路径
img = Image.open(path) # 打开图像
self.imgs.append(pre_trans(img).to(device)) # 预处理图像并添加到列表
self.labels.append(torch.tensor(l_idx).to(device).float()) # 添加标签到列表
def __getitem__(self, idx): # 获取指定索引的数据
img = self.imgs[idx] # 获取图像
label = self.labels[idx] # 获取标签
return img, label # 返回图像和标签
def __len__(self): # 获取数据集的长度
return len(self.imgs) # 返回图像列表的长度
# 定义训练和验证数据集
n = 32 # 定义批次大小
train_path = "data/presidential_doggy_door/train/" # 训练数据路径
train_data = MyDataset(train_path) # 创建训练数据集
train_loader = DataLoader(train_data, batch_size=n, shuffle=True) # 创建训练数据加载器
train_N = len(train_loader.dataset) # 获取训练数据集的大小，用于训练计算准确率
valid_path = "data/presidential_doggy_door/valid/" # 验证数据路径
valid_data = MyDataset(valid_path) # 创建验证数据集
valid_loader = DataLoader(valid_data, batch_size=n) # 创建验证数据加载器
valid_N = len(valid_loader.dataset) # 获取验证数据集的大小，用于验证计算准确率
# 数据增强
IMG_WIDTH, IMG_HEIGHT = (224, 224) # 定义图像宽度和高度
random_trans = transforms.Compose([ # 定义数据增强变换
transforms.RandomRotation(25), # 随机旋转
transforms.RandomResizedCrop((IMG_WIDTH, IMG_HEIGHT), scale=(.8, 1), ratio=(1, 1)), # 随机裁剪
transforms.RandomHorizontalFlip(), # 随机水平翻转
transforms.ColorJitter(brightness=.2, contrast=.2, saturation=.2, hue=.2) # 随机颜色抖动
])
# 定义训练函数
def get_batch_accuracy(output, y, N): # 计算批次准确率的函数
zero_tensor = torch.tensor([0]).to(device) # 创建一个值为 0 的张量
pred = torch.gt(output, zero_tensor) # 获取预测结果
correct = pred.eq(y.view_as(pred)).sum().item() # 计算预测正确的数量
return correct / N # 返回正确预测的比例
# 打印最后一层的梯度
def train(model, check_grad=False): # 定义训练函数
loss = 0 # 初始化损失
accuracy = 0 # 初始化准确率
model.train() # 将模型设置为训练模式
for x, y in train_loader: # 遍历训练数据加载器中的每个批次
output = torch.squeeze(model(random_trans(x))) # 对输入数据进行随机变换并传递给模型
optimizer.zero_grad() # 清零梯度
batch_loss = loss_function(output, y) # 计算损失
batch_loss.backward() # 反向传播
optimizer.step() # 更新模型参数
loss += batch_loss.item() # 累加损失
accuracy += get_batch_accuracy(output, y, train_N) # 累加准确率
if check_grad: # 如果需要检查梯度
print('最后一层梯度:') # 打印最后一层的梯度
for param in model.parameters(): # 遍历模型的参数
print(param.grad) # 打印参数的梯度
print('训练 - 损失: {:.4f} 准确率: {:.4f}'.format(loss, accuracy)) # 打印训练结果
def validate(model): # 定义验证函数
loss = 0 # 初始化损失
accuracy = 0 # 初始化准确率
model.eval() # 将模型设置为评估模式
with torch.no_grad(): # 在不计算梯度的上下文中
for x, y in valid_loader: # 遍历验证数据加载器中的每个批次
output = torch.squeeze(model(x)) # 将输入数据传递给模型
loss += loss_function(output, y.float()).item() # 累加损失
accuracy += get_batch_accuracy(output, y, valid_N) # 累加准确率
print('验证 - 损失: {:.4f} 准确率: {:.4f}'.format(loss, accuracy)) # 打印验证结果
# 开始训练模型
epochs = 3 # 定义训练的轮数，轮数越多，模型的性能可能越好，但因VGG16是很大的模型，训练过长很容易拟合
for epoch in range(epochs): # 遍历每个训练轮次
print('当前: {}'.format(epoch)) # 打印当前轮次
train(my_model, check_grad=False) # 训练模型,不检查梯度
# train(my_model, check_grad=True) # 训练模型，检查梯度
validate(my_model) # 验证模型
# # 保存整个模型到一个新的文件中
# torch.save(my_model.state_dict(), 'my_model.pth') # 保存模型的状态字典到文件
# 预测检测结果
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
# 显示图像
def show_image(image_path):
image = mpimg.imread(image_path)
plt.imshow(image)
# 预测实现函数
def make_prediction(file_path):
show_image(file_path)
image = Image.open(file_path)
image = pre_trans(image).to(device)
image = image.unsqueeze(0)
output = my_model(image)
prediction = output.item()
return prediction
# 判断是否是 Bo
def presidential_doggy_door(image_path):
pred = make_prediction(image_path)
if pred < 0:
print(image_path,"这是 Bo! 欢迎进入!", pred)
else:
print(image_path,"这不是 Bo! 请离开!", pred)
presidential_doggy_door('data/presidential_doggy_door/valid/bo/bo_20.jpg') # 预测结果
presidential_doggy_door('data/presidential_doggy_door/valid/not_bo/121.jpg')