Pytorch入门实战：训练ImageNet数据集

发表于 | 更新于: | 分类于 ，

ImageNet是图像图像领域不可不调的一个数据集。本文记录自己的训练脚本，参考自Pytorch官方的example。

需要掌握的用法如下：

  1. criterion = nn.CrossEntropyLoss().cuda()
    将所有的模型参数移动和buffer移动到gpu上

  2. DataLoader中的参数pin_memory
    将cpu上的tensor数据转移到gpu上时，置为true会加速。
    If you load your samples in the Dataset on CPU and would like to push it during training to the GPU, you can speed up the host to device transfer by enabling pin_memory.

  3. DataLoader中的参数nums_workers
    这个参数必须大于等于0，0表示数据导入在主进程中进行，其他大于0的数表示通过多个进程来导入数据，可以加快数据导入速度。

  4. cudnn.benchmark = True的作用
    大部分情况下，设置这个 flag 可以让内置的 cuDNN 的 auto-tuner 自动寻找最适合当前配置的高效算法，来达到优化运行效率的问题。
    一般来讲，应该遵循以下准则：
    如果网络的输入数据维度或类型上变化不大，设置 torch.backends.cudnn.benchmark = true 可以增加运行效率；
    如果网络的输入数据在每次 iteration 都变化的话，会导致 cnDNN 每次都会去寻找一遍最优配置，这样反而会降低运行效率。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
import argparse
import torch
import torch.optim as optim
import torch.nn as nn
import torch.utils.data.dataloader as dataloader
import torchvision
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms
import os
import warnings
import time
import torch.backends.cudnn as cudnn

model_names = sorted(name for name in models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(models.__dict__[name]))

parser = argparse.ArgumentParser()
parser.add_argument('data', metavar='Dir', help='path to dataset')
parser.add_argument('--arch', '-a',
                    default='resnet18',
                    choices=model_names,
                    help='model architecture: ' +
                         '|'.join(model_names) +
                         ' default(resnet18)')
parser.add_argument('--pretrained', action='store_true', help='use pre-trained model')
parser.add_argument('--learning_rate', '-lr', default='0.1',
                    type=float, help='init learning rate')
parser.add_argument('--start_epoch', default=0, type=int, help='start epoch')
parser.add_argument('--epoch', '-e', default=90, type=int, help='train epochs')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
parser.add_argument('--gpu', action='store_true', help='use gpu')
parser.add_argument('--evaluate', '-eval', action='store_true', help='run in evaluation mode')
parser.add_argument('--batch_size', '-b', default=256, type=int, help='batch size')
parser.add_argument('--print_freq', '-pf', default=30, type=int, help='print info frequent')
parser.add_argument('--weight_decay', '-wd', default=1e-4, type=float, help='weight decay')
parser.add_argument('--workers', '-j', default=4, type=int, help='number of data loading workers(default 4)')

def main():
    global args
    args = parser.parse_args()

    # gpu warning
    if args.gpu:
        warnings.warn("you are using gpu")

    # 定义模型
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=True)
    else:
        print("=>creating model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=False)
    if args.gpu:
        model = model.cuda(0)

    criterion = nn.CrossEntropyLoss().cuda(0)
    optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)

    cudnn.benchmark = True
    # data loader
    normalizer = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    traindir = os.path.join(args.data, 'train')
    transformer_train = transforms.Compose([
        # transforms.Resize(300),
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        normalizer
    ])
    train_dataset = datasets.ImageFolder(
        root=traindir,
        transform=transformer_train,
    )
    train_loader = dataloader.DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.workers,
        pin_memory=True,
    )
    valdir = os.path.join(args.data, 'val')
    transformer_val = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        normalizer
    ])
    val_dataset = datasets.ImageFolder(
        root=valdir,
        transform=transformer_val,
    )
    val_loader = dataloader.DataLoader(
        dataset=val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.workers,
        pin_memory=True,
    )

    if args.evaluate:
        validate(val_loader, model, criterion)
        return

    print("=>start train...")
    for epoch in range(args.start_epoch, args.epoch):
        adjust_learning_rate(optimizer, epoch)
        train(train_loader, model, criterion, optimizer, epoch)
        validate(val_loader, model, criterion)

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def train(train_loader, model, criterion, optimizer, epoch):
    batch_time = AverageMeter()
    losses = AverageMeter()

    # switch to train model
    model.train()

    end = time.time()
    for i, (inputs, labels) in enumerate(train_loader):
        if args.gpu:
            inputs, labels = inputs.cuda(0, non_blocking=True), labels.cuda(0, non_blocking=True)

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        losses.update(loss.item(), inputs.shape[0])

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure time cost
        batch_time.update(time.time() - end)
        end = time.time()

    print(
        'Epoch: {epoch}\t'
        'Loss: {loss.avg:.4f}\t'
        'Batch time cost: {time.sum:.3f}\t'.format(
            epoch=epoch, loss=losses, time=batch_time
        )
    )

def validate(val_loader, model, criterion):
    model.eval()
    accuracy = AverageMeter()
    losses = AverageMeter()

    with torch.no_grad():
        end = time.time()
        for i, (inputs, labels) in enumerate(val_loader):
            if args.gpu:
                inputs, labels = inputs.cuda(non_blocking=True), labels.cuda(non_blocking=True)

            outputs = model(inputs)
            loss = criterion(outputs, labels)
            losses.update(loss.item(), inputs.shape[0])

            _, predicted = torch.max(outputs.data, 1)
            # print((predicted == labels).sum() / inputs.shape[0])
            accuracy.update((predicted == labels).sum() / inputs.shape[0], inputs.shape[0])
        print(
            'Validate loss: {loss.avg:.4f}\t'
            'Validate accuracy: {accuracy.avg:.3f}'.format(
                loss=losses, accuracy=accuracy
            )
        )

def adjust_learning_rate(optimizer, epoch):
    """
    learning rate decay every 30 epochs
    :param optimizer:
    :param epoch:
    :return:
    """
    lr = args.learning_rate * (0.1 ** (epoch // 30))
    for params in optimizer.param_groups:
        params['lr'] = lr


if __name__ == '__main__':
    main()
持续技术分享，您的支持将鼓励我继续创作！