import torch
import torchvision
from torch.autograd import Variable
from torchvision import datasets, transforms
import os # os包集成了一些对文件路径和目录进行操作的类
import matplotlib.pyplot as plt
import time
# 读取数据
data_dir = 'Data' #Data 文件下有train 跟 valid两个文件夹,每个文件下用类别命名,然后放图片
data_transform = {x:transforms.Compose([transforms.Resize([64,64]),
transforms.ToTensor()]) for x in ['train', 'valid']} # 这一步类似预处理
image_datasets = {x:datasets.ImageFolder(root = os.path.join(data_dir,x),
transform = data_transform[x]) for x in ['train', 'valid']} # 这一步相当于读取数据
dataloader = {x:torch.utils.data.DataLoader(dataset = image_datasets[x],
batch_size = 4,
shuffle = True) for x in ['train', 'valid']} # 读取完数据后,对数据进行装载
# 模型搭建
class Models(torch.nn.Module):
def __init__(self):
super(Models, self).__init__()
self.Conv = torch.nn.Sequential(
torch.nn.Conv2d(3, 12, kernel_size=3, stride=1, padding=1),
torch.nn.ReLU(),
torch.nn.Conv2d(12, 24, kernel_size=3, stride=1, padding=1),
torch.nn.ReLU(),
torch.nn.MaxPool2d(kernel_size=2, stride=2))
# torch.nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
# torch.nn.ReLU(),
# torch.nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
# torch.nn.ReLU(),
# torch.nn.MaxPool2d(kernel_size=2, stride=2))
self.Classes = torch.nn.Sequential(
torch.nn.Linear(32 * 32 * 24, 128),
torch.nn.ReLU(),
torch.nn.Dropout(p=0.5),
torch.nn.Linear(128, 3))
def forward(self, inputs):
x = self.Conv(inputs)
x = x.view(-1, 32 * 32 * 24)
x = self.Classes(x)
return x
model = Models()
print(model)
'''
# 保存和加载整个模型
torch.save(model, 'model.pth')
model_1 = torch.load('model.pth')
print(model_1)
# 仅保存和加载模型参数
torch.save(model.state_dict(), 'params.pth')
dic = torch.load('params.pth')
model.load_state_dict(dic)
print(dic)
'''
loss_f = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.00001)
Use_gpu = torch.cuda.is_available()
if Use_gpu:
model = model.cuda()
epoch_n = 100
time_open = time.time()
for epoch in range(epoch_n):
print('epoch {}/{}'.format(epoch, epoch_n - 1))
print('-' * 10)
for phase in ['train', 'valid']:
if phase == 'train':
# # 设置为True,会进行Dropout并使用batch mean和batch var
print('training...')
model.train(True)
else:
# # 设置为False,不会进行Dropout并使用running mean和running var
print('validing...')
model.train(False)
running_loss = 0.0
running_corrects = 0.0
# 输出标号 和对应图片,下标从1开始
for batch, data in enumerate(dataloader[phase], 1):
X, Y = data
# 将数据放在GPU上训练
X, Y = Variable(X).cuda(), Variable(Y).cuda()
# 模型预测概率
y_pred = model(X)
# pred,概率较大值对应的索引值,可看做预测结果,1表示行
_, pred = torch.max(y_pred.data, 1)
# 梯度归零
optimizer.zero_grad()
# 计算损失
loss = loss_f(y_pred, Y)
# 训练 需要反向传播及梯度更新
if phase == 'train':
# 反向传播出现问题
loss.backward()
optimizer.step()
# 损失和
running_loss += loss.data.item()
# 预测正确的图片个数
running_corrects += torch.sum(pred == Y.data)
# 训练时,每500个batch输出一次,训练loss和acc
if batch % 500 == 0 and phase == 'train':
print('batch{},trainLoss:{:.4f},trainAcc:{:.4f}'.format(batch, running_loss / batch,
100 * running_corrects / (4 * batch)))
# 输出每个epoch的loss和acc
epoch_loss = running_loss * 4 / len(image_datasets[phase])
epoch_acc = 100 * running_corrects / len(image_datasets[phase])
print('{} Loss:{:.4f} Acc:{:.4f}%'.format(phase, epoch_loss, epoch_acc))
torch.save(model, 'model.pth')
time_end = time.time() - time_open
print(time_end)
简单图像卷积分类模型(CNN)
最新推荐文章于 2024-07-12 16:28:25 发布