PyTorch学习之路(level1)——训练一个图像分类模型

选取官方链接里面的例子,介绍如何用PyTorch训练一个ResNet模型用于图像分类,代码逻辑非常清晰,基本上和许多深度学习框架的代码思路类似,非常适合初学者想上手PyTorch训练模型(不必每次都跑mnist的demo了)。接下来从个人使用角度加以解释。解释的思路是从数据导入开始到模型训练结束,基本上就是搭积木的方式来写代码。

首先是数据导入部分,这里采用官方写好的torchvision.datasets.ImageFolder接口实现数据导入。这个接口需要你提供图像所在的文件夹,就是下面的data_dir=‘/data’这句,然后对于一个分类问题,这里data_dir目录下一般包括两个文件夹:train和val,每个文件件下面包含N个子文件夹,N是你的分类类别数,且每个子文件夹里存放的就是这个类别的图像。这样torchvision.datasets.ImageFolder就会返回一个列表(比如下面代码中的image_datasets[‘train’]或者image_datasets[‘val]),列表中的每个值都是一个tuple,每个tuple包含图像和标签信息。
————————————

data_dir = '/data'
image_datasets = {x: datasets.ImageFolder(
                    os.path.join(data_dir, x),
                    data_transforms[x]), 
                    for x in ['train', 'val']}

另外这里的data_transforms是一个字典,如下。主要是进行一些图像预处理,比如resize、crop等。实现的时候采用的是torchvision.transforms模块,比如torchvision.transforms.Compose是用来管理所有transforms操作的,torchvision.transforms.RandomSizedCrop是做crop的。需要注意的是对于torchvision.transforms.RandomSizedCrop和transforms.RandomHorizontalFlip()等,输入对象都是PIL Image,也就是用python的PIL库读进来的图像内容,而transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])的作用对象需要是一个Tensor,因此在transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])之前有一个 transforms.ToTensor()就是用来生成Tensor的。另外transforms.Scale(256)其实就是resize操作,目前已经被transforms.Resize类取代了。
————————————————

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomSizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Scale(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

前面torchvision.datasets.ImageFolder只是返回list,list是不能作为模型输入的,因此在PyTorch中需要用另一个类来封装list,那就是:torch.utils.data.DataLoader。torch.utils.data.DataLoader类可以将list类型的输入数据封装成Tensor数据格式,以备模型使用。注意,这里是对图像和标签分别封装成一个Tensor。这里要提到另一个很重要的类:torch.utils.data.Dataset,这是一个抽象类,在pytorch中所有和数据相关的类都要继承这个类来实现。比如前面说的torchvision.datasets.ImageFolder类是这样的,以及这里的torch.util.data.DataLoader类也是这样的。所以当你的数据不是按照一个类别一个文件夹这种方式存储时,你就要自定义一个类来读取数据,自定义的这个类必须继承自torch.utils.data.Dataset这个基类,最后同样用torch.utils.data.DataLoader封装成Tensor。

dataloders = {x: torch.utils.data.DataLoader(image_datasets[x],     
                                            batch_size=4, 
                                            shuffle=True,
                                            num_workers=4) 
                                            for x in ['train', 'val']}

生成dataloaders后再有一步就可以作为模型的输入了,那就是将Tensor数据类型封装成Variable数据类型,来看下面这段代码。dataloaders是一个字典,dataloders[‘train’]存的就是训练的数据,这个for循环就是从dataloders[‘train’]中读取batch_size个数据,batch_size在前面生成dataloaders的时候就设置了。因此这个data里面包含图像数据(inputs)这个Tensor和标签(labels)这个Tensor。然后用torch.autograd.Variable将Tensor封装成模型真正可以用的Variable数据类型。
为什么要封装成Variable呢?在pytorch中,torch.tensor和torch.autograd.Variable是两种比较重要的数据结构,Variable可以看成是tensor的一种包装,其不仅包含了tensor的内容,还包含了梯度等信息,因此在神经网络中常常用Variable数据结构。那么怎么从一个Variable类型中取出tensor呢?也很简单,比如下面封装后的inputs是一个Variable,那么inputs.data就是对应的tensor。

for data in dataloders['train']:
   inputs, labels = data

   if use_gpu:
       inputs = Variable(inputs.cuda())
       labels = Variable(labels.cuda())
   else:
       inputs, labels = Variable(inputs), Variable(labels)

封装好了数据后,就可以作为模型的输入了。所以要先导入你的模型。在PyTorch中已经默认为大家准备了一些常用的网络结构,比如分类中的VGG,ResNet,DenseNet等等,可以用torchvision.models模块来导入。比如用torchvision.models.resnet18(pretrained=True)来导入ResNet18网络,同时指明导入的是已经预训练过的网络。因为预训练网络一般是在1000类的ImageNet数据集上进行的,所以要迁移到你自己数据集的2分类,需要替换最后的全连接层为你所需要的输出。因此下面这三行代码进行的就是用models模块导入resnet18网络,然后获取全连接层的输入channel个数,用这个channel个数和你要做的分类类别数(这里是2)替换原来模型中的全连接层。这样网络结果也准备好。

model = models.resnet18(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)

但是只有网络结构和数据还不足以让代码运行起来,还需要定义损失函数。在PyTorch中采用torch.nn模块来定义网络的所有层,比如卷积、降采样、损失层等等,这里采用交叉熵函数,因此可以这样定义:

criterion = nn.CrossEntropyLoss()

然后你还需要定义优化函数,比如最常见的随机梯度下降,在PyTorch中是通过torch.optim模块来实现的。另外这里虽然写的是SGD,但是因为有momentum,所以是Adam的优化方式。这个类的输入包括需要优化的参数:model.parameters(),学习率,还有Adam相关的momentum参数。现在很多优化方式的默认定义形式就是这样的。

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

然后一般还会定义学习率的变化策略,这里采用的是torch.optim.lr_scheduler模块的StepLR类,表示每隔step_size个epoch就将学习率降为原来的gamma倍。

scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

开始训练

首先训练开始的时候需要先更新下学习率,这是因为我们前面制定了学习率的变化策略,所以在每个epoch开始时都要更新下:

scheduler.step()

然后设置模型状态为训练状态:

model.train(True)

然后先将网络中的所有梯度置0:

optimizer.zero_grad()

然后就是网络的前向传播了:

outputs = model(inputs)

然后将输出的outputs和原来导入的labels作为loss函数的输入就可以得到损失了

loss = criterion(outputs, labels)

输出的outputs也是torch.autograd.Variable格式,得到输出后(网络的全连接层的输出)还希望能到到模型预测该样本属于哪个类别的信息,这里采用torch.max。torch.max()的第一个输入是tensor格式,所以用outputs.data而不是outputs作为输入;第二个参数1是代表dim的意思,也就是取每一行的最大值,其实就是我们常见的取概率最大的那个index;第三个参数loss也是torch.autograd.Variable格式。

 _, preds = torch.max(outputs.data, 1)

计算得到loss后就要回传损失。要注意的是这是在训练的时候才会有的操作,测试时候只有forward过程。

loss.backward()

回传损失过程中会计算梯度,然后需要根据这些梯度更新参数,optimizer.step()就是用来更新参数的。optimizer.step()后,你就可以从optimizer.param_groups[0][‘params’]里面看到各个层的梯度和权值信息。

optimizer.step()

这样一个batch数据的训练就结束了!当你不断重复这样的训练过程,最终就可以达到你想要的结果了。

另外如果你有gpu可用,那么包括你的数据和模型都可以在gpu上操作,这在PyTorch中也非常简单。判断你是否有gpu可以用可以通过下面这行代码,如果有,则use_gpu是true。

use_gpu = torch.cuda.is_available()

参考

PyTorch使用及源码解读
Pytorch-Transfer Learning
Welcome to PyTorch Tutorials
PyTorch学习之路(level1)——训练一个图像分类模型
https://github.com/miraclewkf/ImageClassification-PyTorch/blob/master/level1/train.py

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 156,423评论 4 359
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 66,339评论 1 289
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 106,241评论 0 237
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 43,503评论 0 203
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 51,824评论 3 285
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 40,262评论 1 207
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 31,615评论 2 309
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 30,337评论 0 194
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 33,989评论 1 238
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 30,300评论 2 240
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 31,829评论 1 256
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 28,193评论 2 250
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 32,753评论 3 230
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 25,970评论 0 8
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 26,708评论 0 192
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 35,295评论 2 267
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 35,207评论 2 258

推荐阅读更多精彩内容