1.安装
https://pytorch.org/ 官网,根据对应情况执行命令行,就完了
2.数据预处理
from PIL import Image #读图
from torchvision import transforms, datasets #导入图像数据处理,转换库
import os
import torch
def readImg(path):
'''
用于替代ImageFolder的默认读取图片函数,以读取单通道图片
'''
return Image.open(path)
def ImageDataset(args):
data_transforms = { #数据增强obj
'train': transforms.Compose([ #组合转换
transforms.RandomCrop(RANDOMCROP_SIZE), #随机裁剪
transforms.RandomHorizontalFlip(), #随机水平翻转
transforms.ToTensor(), #转tensor 变float变tensor
transforms.Normalize([0.5], [0.5]) #改均值
]),
'test': transforms.Compose([#测试时要和train对应,有些增强不需要
transforms.CenterCrop(RANDOMCROP_SIZE),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
]),
}
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), #实用读分类图库
data_transforms[x] #数据增强
# , loader=readImg #单通道图像使用
)
for x in ['train', 'test']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=args.batch_size, shuffle=(x == 'train'),
num_workers=args.num_workers)
for x in ['train', 'test']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'test']}
class_names = image_datasets['train'].classes
return dataloaders, dataset_sizes, class_names
3.网络模型保存和加载
(1) 仅仅保存和加载模型参数
torch.save(the_model.state_dict(), PATH)
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))
(2) 保存和加载整个模型
torch.save(the_model, PATH)
#切记! 对应网络class 需要在加载之前已经定义
the_model = torch.load(PATH)
#使用GPU时
model = TheModelClass().cuda()
model = nn.DataParallel(model).cuda()
model.load_state_dict(torch.load(best_model_path))
model.eval()
# 预测图像的类别
output = model(input)
4.网络构建
import torch
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
# 四个卷积层用于提取特征
# 3 input channel image 160x160, 8 output channel image 80x80
self.conv = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=8, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
)
# 3 input channel image 80x80, 16 output channel image 40x40
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
)
# 8 input channel image 40x40, 32 output channel image 20x20
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
)
# 16 input channel image 20x20, 64 output channel image 10x10
self.conv3 = nn.Sequential(
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
)
# 分类
self.classifier = nn.Sequential(
nn.Linear(64 * 10 * 10, 6)
)
def forward(self, x):
x = self.conv(x)
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = x.view(-1, 64 * 10 * 10)
x = self.classifier(x)
return x
5.训练过程
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
model = SimpleNet()
#是否使用gpu
if use_gpu:
model = torch.nn.DataParallel(model)
model.to(torch.device('cuda'))
else:
model.to(torch.device('cpu'))
确定loss Function
criterion = nn.CrossEntropyLoss()
确定优化方法
optimizer_ft = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.00004)
optimizer_ft = optim.Adam(model.parameters(), lr=args.lr)
学习率衰减
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=1, gamma=0.98)
device = torch.device('cuda' if use_gpu else 'cpu')
#test,log之类的没写,可以自己加,复制是不能直接跑的,只是介绍流程
for epoch in need_epoch:
exp_lr_scheduler (epoch)
model.train() #训练模式,batchNormal和dropout会跑,model.eval()时不跑
# 在多个batch上依次处理数据(Iterate over data)
for i, (inputs, labels) in enumerate(dataloders[phase]):
inputs = inputs.to(device)
labels = labels.to(device)
# 梯度置零(zero the parameter gradients)
optimizer.zero_grad()
# 前向传播(forward)
# 训练模式下才记录梯度以进行反向传播(track history if only in train)
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
# 训练模式下进行反向传播与梯度下降(backward + optimize only if in training phase)
loss.backward()
optimizer.step()
# 统计损失和准确率(statistics)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
batch_loss = running_loss / (i * args.batch_size + inputs.size(0))
batch_acc = running_corrects.double() / (i * args.batch_size + inputs.size(0))
epoch_loss = running_loss / dataset_sizes
epoch_acc = running_corrects.double() / dataset_sizes
# 深拷贝模型(deep copy the model)
if phase == 'test' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
6.tensorboard
安装 pip install tensorboard==2.0.0 (固定安装2.0版本,因为torch版本未到2.1,后面如果上去了可以自己改 )
# from tensorboardX import SummaryWriter
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir='log') #存储log的位置
writer.add_scalar(phase + '/Loss', epoch_loss, epoch) #数值型(表title,y轴数值,x轴数值)
writer.add_image(phase + '/Loss', image, epoch) #数值型(表title,图像值,x轴数值)
writer.close()
查看tensorboard:在存储log的位置运行命令行 tensorboard --logdir runs
注:有参考网上一个教程,具体链接找不到了。。。。