首发于深度学习

pytorch训练自己图像分类数据集

2020年6月更新,调整代码的组成结构,使其更加通用,方便使用,更加整洁,添加模型蒸馏的方法。

希望对初学者能有些帮助,欢迎star

lxztju/pytorch_classificationgithub.com图标

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

https://github.com/lxztju/pytorch_classificationgithub.com



近期无意间看到华为云的垃圾分类大赛,看了官方的baseline代码,发现是使用keras写的,自己没有用过keras,看不懂代码,因而就用pytorch写了代码,利用densenet模型进行训练,完整的走了一遍这种图像分类的过程,从数据集的格式到模型的训练、测试与单张图片的预测,到最后部署在华为云的modelarts上,整个图像分类工作的流程

下载完整代码,请移步github,如果有用的话,不要忘了‘star’

lxztju/densenet-pytorchgithub.com图标

这里记录一下自己的工作,同时也给刚入门深度学习、刚开始学习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上,这里就先不叙述了

ModelArts平台运行PyTorch模型_ModelArts_EI企业智能_热门产品_华为云bbs.huaweicloud.com图标

编辑于 06-28

文章被以下专栏收录

    简单随笔记录自己深度学习的学习过程,帮助自己回忆的同时,希望能够帮助到有需要的同学少踩点坑