Skip to content

Latest commit

 

History

History

Folders and files

NameName
Last commit message
Last commit date

parent directory

..
 
 
 
 
 
 
 
 
 
 
 
 

ResNet Series

本目录包含了采用MegEngine实现的经典ResNet网络结构,同时提供了在ImageNet训练集上的完整训练和测试代码。

model.py中定义了一些常见的网络结构:resnet18, resnet34, resnet50, resnet101, resnet152等.

目前我们提供了部分在ImageNet上的预训练模型(见下表),各个网络结构在ImageNet验证集上的测试结果如下:

模型 top1 acc top5 acc
ResNet18 70.312 89.430
ResNet34 73.960 91.630
ResNet50 76.254 93.056
ResNet101 77.944 93.844
ResNet152 78.582 94.130
ResNeXt50 32x4d 77.592 93.644
ResNeXt101 32x8d 79.520 94.586

用户可以通过megengine.hub直接加载本目录下定义好的模型,例如:

import megengine.hub

# 只加载网络结构
resnet18 = megengine.hub.load("megengine/models", "resnet18")
# 加载网络结构和预训练权重
resnet18 = megengine.hub.load("megengine/models", "resnet18", pretrained=True)

安装和环境配置

在开始运行本目录下的代码之前,请确保按照README进行了正确的环境配置。

如何训练

在开始训练前,请确保已经下载解压好ImageNet数据集,并放在合适的目录下,准备好的数据集的目录结构如下所示:

/path/to/imagenet
    train
         n01440764
              xxx.jpg
              ...
         n01443537
              xxx.jpg
              ...
         ...
    val
         n01440764
              xxx.jpg
              ...
         n01443537
              xxx.jpg
              ...
         ...

准备好数据集后,可以运行以下命令开始训练:

python3 train.py --dataset-dir=/path/to/imagenet

train.py提供了灵活的命令行选项,包括:

  • --data, ImageNet数据集的根目录,默认/data/datasets/imagenet;
  • --arch, 需要训练的网络结构,默认resnet50
  • --batch-size,训练时每张卡采用的batch size, 默认64;
  • --ngpus, 训练时每个节点采用的gpu数量,默认None,即使用全部gpu;当使用多张gpu时,将自动切换为分布式训练模式;
  • --save, 模型以及log存储的目录,默认output;
  • --learning-rate, 训练时的初始学习率,默认0.025,在分布式训练下,实际学习率等于初始学习率乘以总gpu数;
  • --epochs, 训练多少个epoch,默认90;

例如,可以通过以下命令在2块GPU上以64的batch大小训练一个resnet50的模型:

python3 train.py --data /path/to/imagenet \
                 --arch resnet50 \
                 --batch-size 64 \
                 --learning-rate 0.025 \
                 --ngpus 2 \
                 --save /path/to/save_dir

更多详细的介绍可以通过运行python3 train.py --help查看。

如何测试

在训练的过程中,可以通过如下命令测试模型在ImageNet验证集的性能:

python3 test.py --data=/path/to/imagenet --arch resnet50 --model /path/to/model --ngpus 1

test.py的命令行选项如下:

  • --data,ImageNet数据集的根目录,默认/data/datasets/imagenet
  • --arch, 需要测试的网络结构,默认resnet50
  • --model, 需要测试的模型,默认使用官方预训练模型;
  • --ngpus, 用于测试的gpu数量,默认None

更多详细介绍可以通过运行python3 test.py --help查看。

如何使用

模型训练好之后,可以通过如下命令测试单张图片:

python3 inference.py --model /path/to/model --image /path/to/image.jpg

使用默认的测试图片和默认的resnet18模型,将输出如下结果:

0: class = lynx                 with probability = 25.2 %
1: class = Siamese_cat          with probability = 12.3 %
2: class = Egyptian_cat         with probability =  8.7 %
3: class = Persian_cat          with probability =  8.3 %
4: class = tabby                with probability =  6.5 %

inference.py的命令行选项如下:

  • --arch, 需要使用的网络结构,默认resnet18
  • --model, 训练好的模型权重地址,默认使用官方预训练的resnet18模型;
  • --image, 用于测试的图片;

参考文献