1、resnet训练代码如下：
import torch
from torchvision import datasets, models, transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import time
import numpy as np
import matplotlib.pyplot as plt
import os
from tqdm import tqdm
# 一、建立数据集
# animals-6
# --train
# |--dog
# |--cat
# ...
# --valid
# |--dog
# |--cat
# ...
# --test
# |--dog
# |--cat
# ...
# 我的数据集中 train 中每个类别60张图片，valid 中每个类别 10 张图片，test 中每个类别几张到几十张不等，一共 6 个类别。
# 二、数据增强
# 建好的数据集在输入网络之前先进行数据增强，包括随机 resize 裁剪到 256 x 256，随机旋转，随机水平翻转，中心裁剪到 224 x 224，转化成 Tensor，正规化等。
image_transforms = {
'train': transforms.Compose([
transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),
transforms.RandomRotation(degrees=15),
transforms.RandomHorizontalFlip(),
transforms.CenterCrop(size=224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
]),
'valid': transforms.Compose([
transforms.Resize(size=256),
transforms.CenterCrop(size=224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
}
# 三、加载数据
# torchvision.transforms包DataLoader是 Pytorch 重要的特性，它们使得数据增加和加载数据变得非常简单。
# 使用 DataLoader 加载数据的时候就会将之前定义的数据 transform 就会应用的数据上了。
dataset = '/home/user/workspace/resnet/car1_data'
path = 'resnet50'
train_directory = os.path.join(dataset, 'train')
valid_directory = os.path.join(dataset, 'val')
batch_size = 32
num_classes = 2
print(train_directory)
data = {
'train': datasets.ImageFolder(root=train_directory, transform=image_transforms['train']),
'valid': datasets.ImageFolder(root=valid_directory, transform=image_transforms['valid'])
}
train_data_size = len(data['train'])
valid_data_size = len(data['valid'])
train_data = DataLoader(data['train'], batch_size=batch_size, shuffle=True, num_workers=8)
valid_data = DataLoader(data['valid'], batch_size=batch_size, shuffle=True, num_workers=8)
print(train_data_size, valid_data_size)
# 四、迁移学习
# 这里使用ResNet-50的预训练模型。
resnet50 = models.resnet50(pretrained=True)
# 在PyTorch中加载模型时，所有参数的‘requires_grad’字段默认设置为true。这意味着对参数值的每一次更改都将被存储，以便在用于训练的反向传播图中使用。
# 这增加了内存需求。由于预训练的模型中的大多数参数已经训练好了，因此将requires_grad字段重置为false。
for param in resnet50.parameters():
param.requires_grad = False
# 为了适应自己的数据集，将ResNet-50的最后一层替换为，将原来最后一个全连接层的输入喂给一个有256个输出单元的线性层，接着再连接ReLU层和Dropout层，然后是256 x 6的线性层，输出为6通道的softmax层。
fc_inputs = resnet50.fc.in_features
resnet50.fc = nn.Sequential(
nn.Linear(fc_inputs, 256),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(256, 2),
nn.LogSoftmax(dim=1)
)
# 用GPU进行训练。
resnet50 = resnet50.to('cuda:0')
# 定义损失函数和优化器。
loss_func = nn.NLLLoss()
optimizer = optim.Adam(resnet50.parameters())
# 五、训练
def train_and_valid(model, loss_function, optimizer, epochs=25):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
history = []
best_acc = 0.0
best_epoch = 0
for epoch in range(epochs):
epoch_start = time.time()
print("Epoch: {}/{}".format(epoch+1, epochs))
model.train()
train_loss = 0.0
train_acc = 0.0
valid_loss = 0.0
valid_acc = 0.0
for i, (inputs, labels) in enumerate(tqdm(train_data)):
inputs = inputs.to(device)
labels = labels.to(device)
#因为这里梯度是累加的，所以每次记得清零
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
train_loss += loss.item() * inputs.size(0)
ret, predictions = torch.max(outputs.data, 1)
correct_counts = predictions.eq(labels.data.view_as(predictions))
acc = torch.mean(correct_counts.type(torch.FloatTensor))
train_acc += acc.item() * inputs.size(0)
with torch.no_grad():
model.eval()
for j, (inputs, labels) in enumerate(tqdm(valid_data)):
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
loss = loss_function(outputs, labels)
valid_loss += loss.item() * inputs.size(0)
ret, predictions = torch.max(outputs.data, 1)
correct_counts = predictions.eq(labels.data.view_as(predictions))
acc = torch.mean(correct_counts.type(torch.FloatTensor))
valid_acc += acc.item() * inputs.size(0)
avg_train_loss = train_loss/train_data_size
avg_train_acc = train_acc/train_data_size
avg_valid_loss = valid_loss/valid_data_size
avg_valid_acc = valid_acc/valid_data_size
history.append([avg_train_loss, avg_valid_loss, avg_train_acc, avg_valid_acc])
if best_acc < avg_valid_acc:
best_acc = avg_valid_acc
best_epoch = epoch + 1
epoch_end = time.time()
print("Epoch: {:03d}, Training: Loss: {:.4f}, Accuracy: {:.4f}%, \n\t\tValidation: Loss: {:.4f}, Accuracy: {:.4f}%, Time: {:.4f}s".format(
epoch+1, avg_valid_loss, avg_train_acc*100, avg_valid_loss, avg_valid_acc*100, epoch_end-epoch_start
))
print("Best Accuracy for validation : {:.4f} at epoch {:03d}".format(best_acc, best_epoch))
#torch.save(model, 'models/'+dataset+'_model_'+str(epoch+1)+'.pt')\
torch.save(model.state_dict(), 'models/'+ path +'_model_'+str(epoch+1)+'.pt')
return model, history
num_epochs = 30
trained_model, history = train_and_valid(resnet50, loss_func, optimizer, num_epochs)
torch.save(history.state_dict(), 'models/'+ path +'_history.pt')
history = np.array(history)
plt.plot(history[:, 0:2])
plt.legend(['Tr Loss', 'Val Loss'])
plt.xlabel('Epoch Number')
plt.ylabel('Loss')
plt.ylim(0, 1)
plt.savefig(dataset+'_loss_curve.png')
plt.show()
plt.plot(history[:, 2:4])
plt.legend(['Tr Accuracy', 'Val Accuracy'])
plt.xlabel('Epoch Number')
plt.ylabel('Accuracy')
plt.ylim(0, 1)
plt.savefig(dataset+'_accuracy_curve.png')
plt.show()
2、resnet50的pt模型验证代码
# -*- coding: utf-8 -*-
import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import torch.nn as nn
model = models.resnet50(pretrained=False) # 设置pretrained=False，不使用预训练的权重
#num_ftrs = model.fc.in_features
#model.fc = torch.nn.Linear(num_ftrs, num_classes) # 替换全连接层
fc_inputs = model.fc.in_features
model.fc = nn.Sequential(
nn.Linear(fc_inputs, 256),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(256, 2),
nn.LogSoftmax(dim=1)
)
model.load_state_dict(torch.load("/mnt/hgfs/shell/resnet50_model_28.pt"))
model.eval() # 设置模型为评估模式
transform = 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]),
])
image = Image.open("/home/user/workspace/resnet/car1_data/test/bat.jpg")
print("load image")
image = transform(image).unsqueeze(0) # 增加批处理维度
print("image process")
with torch.no_grad():
outputs = model(image)
_, predicted = torch.max(outputs, 1)
print('Predicted class:', predicted.item())
_, predicted = torch.max(outputs, 1)
print('Predicted class:', predicted.item())3、resnet50 pt模型转换为onnx格式代码：
import os
import torch
import torch.onnx
from torch.autograd import Variable
from torchvision import models
import torch.nn as nn
from PIL import Image
def get_pytorch_onnx_model(original_model):
# define the directory for further converted model save
onnx_model_path = "/home/ck/work/resnet"
# define the name of further converted model
onnx_model_name = "resnet50_20.onnx"
# create directory for further converted model
os.makedirs(onnx_model_path, exist_ok=True)
# get full path to the converted model
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
# generate model input
generated_input = Variable(
torch.randn(1, 3, 224, 224)
)
# model export into ONNX format
torch.onnx.export(
original_model,
generated_input,
full_model_path,
verbose=True,
input_names=["input"],
output_names=["output"],
opset_version=11
)
return full_model_path
def main():
# initialize PyTorch MobileNetV2
#original_model = models.resnet50(pretrained=True)
model = models.resnet50(pretrained=False) # 设置pretrained=False，不使用预训练的权重
fc_inputs = model.fc.in_features
model.fc = nn.Sequential(nn.Linear(fc_inputs, 256),nn.ReLU(),nn.Dropout(0.4),nn.Linear(256, 2),nn.LogSoftmax(dim=1))
model_path = "/mnt/hgfs/shell/resnet50_model_20.pt"
checkpoint = torch.load(model_path,map_location=torch.device('cpu'))
model.load_state_dict(checkpoint)
# get the path to the converted into ONNX PyTorch model
full_model_path = get_pytorch_onnx_model(model)
print("PyTorch resnet50 model was successfully converted: ", full_model_path)
if __name__ == "__main__":
main()4、推理onnx格式resnet50模型
import onnxruntime
import numpy as np
from torchvision import datasets, models, transforms
from PIL import Image
import torch.nn as nn
import torch
# 定义ONNX模型路径
onnx_model_path = "/home/ck/work/resnet/resnet50_20.onnx"
# 加载ONNX模型
ort_session = onnxruntime.InferenceSession(onnx_model_path)
# 准备输入数据，这里使用一个随机数组模拟
# 在实际应用中，你应该使用适当的数据预处理
#input_data = np.random.rand(1, 3, 224, 224).astype(np.float32)
transform = 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]),
])
image = Image.open("/home/ck/work/resnet/data/test/front0.jpg")
image = transform(image).unsqueeze(0) # 增加批处理维度
# 将 PyTorch 张量转换为 NumPy 数组
input_data = image.detach().numpy()
# 推理
# 注意：'input'是ONNX模型输入节点的名称，你可能需要根据你的模型具体情况进行修改
outputs_np = ort_session.run(None, {'input': input_data})
print(outputs_np)
outputs = outputs_np[0]
print(outputs)
outputs_tensor = torch.tensor(outputs)
max_value, max_indices = torch.max(outputs_tensor,1)
# 输出推理结果
#print(max_value)
print(max_indices.item())
1310
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
