pytorch训练自己图像分类数据集
2020年6月更新，调整代码的组成结构，使其更加通用，方便使用，更加整洁，添加模型蒸馏的方法。
希望对初学者能有些帮助，欢迎star
2020年4月26更新：cnn svm knn pytorch
最近实现了利用cnn提取特征，然后利用svm或者knn，随机森林等分类器实现分类，在一些分类任务中效果会更好，代码已经在github仓库中更新。
代码实现主要功能是，cnn训练完成后，去掉全连接层，然后将提取到的训练集的特征保存为pkl文件，然后利用这些特征训练svm或者knn等分类器，保存分类器，并读取预测
2020年4月更新
之前利用pytorch写的代码，是自己刚开始使用pytorch学习时，自己练手写的，最近慢慢完善了一个pytorch实现图像分类的代码。
实现主要功能：
- 基础功能利用pytorch实现图像分类
- 包含带有warmup的cosine学习率优化
- 多模型融合预测，加权与投票融合
- 利用flask实现模型云端api部署
查看完整代码，请移步github，如果有用的话，不要忘记star
近期无意间看到华为云的垃圾分类大赛，看了官方的baseline代码，发现是使用keras写的，自己没有用过keras，看不懂代码，因而就用pytorch写了代码，利用densenet模型进行训练，完整的走了一遍这种图像分类的过程，从数据集的格式到模型的训练、测试与单张图片的预测，到最后部署在华为云的modelarts上，整个图像分类工作的流程
下载完整代码，请移步github，如果有用的话，不要忘了‘star’
这里记录一下自己的工作，同时也给刚入门深度学习、刚开始学习pytorch的同学一个参考，给大家一个相对简单的实现过程，简单的代码实现，其中也介绍了很多注意的要点
例如：1、采用ImageFolder与dataloader制作提取数据集的提取器;
2、pytorch 0.4之后的版本中Tensor与Variable合并，Variable已经不再使用，Volatile = True已经不再使用，由torch.no_grad()代替
3、由于在分类任务中，一般需要读取预训练权重进行预处理，加快网络的收敛。但是，我们一般的分类的类别数目不尽相同，因此许啊哟读入前边几层，最后一层的参数利用初始化函数进行初始化，这里介绍了读取部分预训练模型参数的方法，可以很好的理解pytorch保存的模型为字典格式
4、介绍多gpu并行训练的方式
5、介绍多gpu训练的模型保存参数与单gpu训练保存的模型参数的不同点，以及加载多gpu模型时需要对模型的键进行处理，去掉‘module’
6、在进行预测评估时，使用eval()函数
这些在这篇文章中都能找到答案，相信读完这篇文章你也可以写一个多gpu并行运算的图像分类算法
1、准备自己的数据集
对于图像分类任务，将每一类的照片放置在同一个文件夹下，然后将所有包含各种图像的文件夹放置在同一个文件夹下，例如：
我们有10类图像，分别放在10个文件夹下，命名为1,2,3,4,5,6,7,8,9,10。然后将这十个文件夹放置在同一个文件夹下，这个文件夹命名为train_data，
得到数据集之后，利用github中代码文件load_data.py来创建数据提取器，这返回一个迭代器
利用torchvision.datasets中的ImageFolder类直接进行封装，这个类直接将每一个文件夹下的图片与类别对应，返回结果为迭代器。 然后将这个迭代器传入dataloader，按每个batch提取数据
主要代码段理解：
##
train_transforms = transforms.Compose([ transforms.RandomRotation(10), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((.5, .5, .5), (.5, .5, .5))])
train_dir = cfg.TRAIN_DATASET_DIR
train_datasets = datasets.ImageFolder(train_dir, transform=train_transforms)
train_dataloader = torch.utils.data.DataLoader(train_datasets, batch_size=batch_size, shuffle=True, num_workers=2)2、densenet模型文件
参考torchvision中的models
3、训练代码
github代码中的train.py
##读入下载的预训练的权重字典
state_dict = torch.load(pretrained_path)
###去掉全链接层的权重，
#由于我们一般不会直接使用，imagenet的1000类，因此，我们需要更换网络最后的全链接层
#因此我们需要将前边几层的参数保存，最后一层重新初始化
#定义一个新的字典，将原始的参数字典，对应保存与更改
from collections import OrderedDict
new_state_dict = OrderedDict()
for k,v in state_dict.items():
# print(k) #打印预训练模型的键，发现与网络定义的键有一定的差别，因而需要将键值进行对应的更改，将键值分别对应打印出来就可以看出不同，根据不同进行修改
#torchvision中的网络定义，采用了正则表达式，来更改键值，因为这里简单，没有再去构建正则表达式
# 直接利用if语句筛选不一致的键
###修正键值的不对应
if k.split('.')[0] == 'features' and (len(k.split('.')))>4:
k = k.split('.')[0]+'.'+k.split('.')[1]+'.'+k.split('.')[2]+'.'+k.split('.')[-3] + k.split('.')[-2] +'.'+k.split('.')[-1]
# print(k)
else:
pass
##最后一层的全连接层，进行初始化
if k.split('.')[0] == 'classifier':
if k.split('.')[-1] == 'weights':
v = nn.init.kaiming_normal(model.state_dict()[k], mode='fan_out')
else:
model.state_dict()[k][...] = 0.0
v = model.state_dict()[k][...]
else:
pass
##得到新的与定义网络对应的预训练参数
new_state_dict[k] = v
##导入网络参数
model.load_state_dict(new_state_dict)
##进行多gpu的并行计算
if args.ngpu:
model = nn.DataParallel(model,device_ids=list(range(args.ngpu)))
print("initialize the network done")
##定义优化器与损失函数optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
loss_func = nn.CrossEntropyLoss()训练过程
##训练max——epoch个epoch
for epoch in range(args.max_epoch):
model.train() ##在进行训练时加上train()，测试时加上eval()
##在测试时加上eval()会将BN与Dropout的进行固定
batch = 0
for batch_images, batch_labels in train_dataloader:
# print(batch_labels)
# print(torch.cuda.is_available())
average_loss = 0
train_acc = 0
##在pytorch0.4之后将Variable 与tensor进行合并，所以这里不需要进行Variable封装
if torch.cuda.is_available():
batch_images, batch_labels = batch_images.cuda(),batch_labels.cuda()
out = model(batch_images)
loss = loss_func(out,batch_labels)
# print(loss)
average_loss = loss
prediction = torch.max(out,1)[1]
# print(prediction)
train_correct = (prediction == batch_labels).sum()
##这里得到的train_correct是一个longtensor型，需要转换为float
# print(train_correct.type())
train_acc = (train_correct.float()) / args.batch_size
# print(train_acc.type())
optimizer.zero_grad() #清空梯度信息，否则在每次进行反向传播时都会累加
loss.backward() #loss反向传播
optimizer.step() ##梯度更新
batch+=1
print("Epoch: %d/%d || batch:%d/%d average_loss: %.3f || train_acc: %.2f"
%(epoch, args.max_epoch, batch, max_batch, average_loss, train_acc))
##每10epoch保存一次模型
if epoch%10 ==0 and epoch>0:
torch.save(model.state_dict(), args.save_folder+'/'+'dense169'+'_'+str(epoch)+'.pth')
4、测试网络模型效果
github代码中eval.py
###读取网络模型的键值对
trained_model = cfg.TRAINED_MODEL
state_dict = torch.load(trained_model)
# create new OrderedDict that does not contain `module.`
##由于之前的模型是在多gpu上训练的，因而保存的模型参数，键前边有‘module’，需要去掉，和训练模型一样构建新的字典
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
print('Finished loading model!')
##进行模型测试时，eval（）会固定下BN与Dropout的参数
model.eval()
eval_acc = 0.0
for batch_images, batch_labels in val_dataloader:
# print(batch_labels)
with torch.no_grad():
if torch.cuda.is_available():
batch_images, batch_labels = batch_images.cuda(), batch_labels.cuda()
##在pytorch0.4的版本之前，使用Variable封装，里边采用volatile=True放置进行反传训练
#在0.4之后，官方推荐torch.no_grad()，Variable API已经被弃用
out = model(batch_images)
prediction = torch.max(out, 1)[1]
num_correct = (prediction == batch_labels).sum()
eval_acc += num_correct
print(eval_acc)
print(' Accuracy of batch : {:.6f}'.format((eval_acc.float()) / (len(val_datasets))))
5、单张图片结果预测
github代码中的predict.py
##labels_to_classes为与输出结果对应的字典，例如：{'0','flower'}
labels2classes = cfg.labels_to_classes
###读取网络模型的键值对
trained_model = cfg.TRAINED_MODEL
state_dict = torch.load(trained_model)
# create new OrderedDict that does not contain `module.`
##由于之前的模型是在多gpu上训练的，因而保存的模型参数，键前边有‘module’，需要去掉，和训练模型一样构建新的字典
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
print('Finished loading model!')
##进行模型测试时，eval（）会固定下BN与Dropout的参数
model.eval()
with torch.no_grad():
if torch.cuda.is_available():
batch_images, batch_labels = batch_images.cuda(), batch_labels.cuda()
##在pytorch0.4的版本之前，使用Variable封装，里边采用volatile=True放置进行反传训练
#在0.4之后，官方推荐torch.no_grad()，Variable PI已经被弃用
out = model(batch_images)
prediction = torch.max(out, 1)[1]
print(labels2classes[str(prediction)]6、华为云modelarts部署
deployment中的代码主要参考这个，主要应用于将模型部署到华为云modelarts上，这里就先不叙述了
