02. Train Image Classification with Auto Estimator

This tutorial goes through the basic steps of using GluonCV auto estimator to train an image classifier with custom hyper-parameters.

Train with default configurations

from gluoncv.auto.estimators import ImageClassificationEstimator

In this tutorial, we use a small sample dataset

train, _, test = ImageClassificationEstimator.Dataset.from_folders(
    'https://autogluon.s3.amazonaws.com/datasets/shopee-iet.zip')
train, val, _ = train.random_split(val_size=0.1, test_size=0)

Out:

data/
├── test/
└── train/

Create an estimator with default configurations. We only change the number of GPUs to reflect the hardware constraint.

Note that you may still launch training if no gpu is available(with {‘gpus’: []}) but be prepared that this is painfully slow and not even possible to finish in case the dataset isn’t tiny.

We recommend that you use at least one nvidia gpu with more than 6GB free GPU memory.

classifier = ImageClassificationEstimator(
    {'gpus': [0], 'train': {'batch_size': 16, 'epochs': 2}})

run fit on train/validation data

classifier.fit(train, val)

Out:

Downloading /root/.mxnet/models/resnet50_v1-cc729d95.zip from https://apache-mxnet.s3-accelerate.dualstack.amazonaws.com/gluon/models/resnet50_v1-cc729d95.zip...

  0%|          | 0/57421 [00:00<?, ?KB/s]
  0%|          | 99/57421 [00:00<01:12, 793.84KB/s]
  1%|          | 510/57421 [00:00<00:25, 2253.13KB/s]
  4%|3         | 2189/57421 [00:00<00:07, 7331.41KB/s]
 14%|#3        | 7815/57421 [00:00<00:02, 24063.87KB/s]
 25%|##4       | 14265/57421 [00:00<00:01, 37334.23KB/s]
 40%|###9      | 22933/57421 [00:00<00:00, 53084.51KB/s]
 53%|#####3    | 30480/57421 [00:00<00:00, 60094.47KB/s]
 67%|######7   | 38483/57421 [00:00<00:00, 65610.12KB/s]
 82%|########1 | 46854/57421 [00:00<00:00, 71114.66KB/s]
 95%|#########5| 54600/57421 [00:01<00:00, 73035.76KB/s]
100%|##########| 57421/57421 [00:01<00:00, 51737.76KB/s]

Evaluate the final model on test set

eval_result = classifier.evaluate(test)
print("Top1/Top5 on test dataset: {}".format(eval_result))

Out:

Top1/Top5 on test dataset: (0.8875, 1.0)

save to/from disk to be used later

classifier.save('classifier.pkl')
classifier2 = ImageClassificationEstimator.load('classifier.pkl')

run prediction on test images

pred = classifier2.predict(test.iloc[0]['image'])
print('GroundTruth:', test.iloc[0])
print('prediction:', pred)

Out:

GroundTruth: image    /root/.gluoncv/datasets/shopee-iet/data/test/B...
label                                                    0
Name: 0, dtype: object
prediction:               class     score  id
0         BabyPants  0.587603   0
1         BabyShirt  0.388625   1
2   womenchiffontop  0.014393   3
3  womencasualshoes  0.009378   2

Customize your training configurations

You may modify configurations to customize your training, supported fields:

print(ImageClassificationEstimator._default_cfg)

Out:

ImageClassificationCfg(img_cls=ImageClassification(model=resnet50_v1, use_pretrained=True, use_gn=False, batch_norm=False, use_se=False, last_gamma=False), train=TrainCfg(pretrained_base=True, batch_size=128, epochs=10, lr=0.1, lr_decay=0.1, lr_decay_period=0, lr_decay_epoch=40, 60, lr_mode=step, warmup_lr=0.0, warmup_epochs=0, num_training_samples=1281167, num_workers=4, wd=0.0001, momentum=0.9, teacher=None, hard_weight=0.5, dtype=float32, input_size=224, crop_ratio=0.875, use_rec=False, rec_train=~/.mxnet/datasets/imagenet/rec/train.rec, rec_train_idx=~/.mxnet/datasets/imagenet/rec/train.idx, rec_val=~/.mxnet/datasets/imagenet/rec/val.rec, rec_val_idx=~/.mxnet/datasets/imagenet/rec/val.idx, data_dir=~/.mxnet/datasets/imagenet, mixup=False, no_wd=False, label_smoothing=False, temperature=20, resume_epoch=0, mixup_alpha=0.2, mixup_off_epoch=0, log_interval=50, mode=, start_epoch=0, transfer_lr_mult=0.01, output_lr_mult=0.1, early_stop_patience=-1, early_stop_min_delta=0.001, early_stop_baseline=0.0, early_stop_max_value=1.0), valid=ValidCfg(batch_size=128, num_workers=4), gpus=(0,))

For example, we could change the learning rate and batch size

new_classifier = ImageClassificationEstimator({'gpus': [0], 'train': {'batch_size': 16, 'lr': 0.01}})

A more natural format for modifying individual hyperparameter is to edit the yaml file saved automatically in self._logdir, here we just show an example of how to copy/edit and load the modified configuration file back to the estimator in python

You may edit the yaml file directly with a text editor

import shutil
import os
config_name = 'config.yaml'
shutil.copyfile(os.path.join(classifier2._logdir, config_name), os.path.join('.', config_name))
cfg = open(config_name).read()
print(cfg)

Out:

# ImageClassificationCfg
gpus: !!python/tuple
- 0
img_cls:
  batch_norm: false
  last_gamma: false
  model: resnet50_v1
  use_gn: false
  use_pretrained: true
  use_se: false
train:
  batch_size: 16
  crop_ratio: 0.875
  data_dir: ~/.mxnet/datasets/imagenet
  dtype: float32
  early_stop_baseline: 0.0
  early_stop_max_value: 1.0
  early_stop_min_delta: 0.001
  early_stop_patience: -1
  epochs: 10
  hard_weight: 0.5
  input_size: 224
  label_smoothing: false
  log_interval: 50
  lr: 0.01
  lr_decay: 0.1
  lr_decay_epoch: 40, 60
  lr_decay_period: 0
  lr_mode: step
  mixup: false
  mixup_alpha: 0.2
  mixup_off_epoch: 0
  mode: ''
  momentum: 0.9
  no_wd: false
  num_training_samples: 1281167
  num_workers: 4
  output_lr_mult: 0.1
  pretrained_base: true
  rec_train: ~/.mxnet/datasets/imagenet/rec/train.rec
  rec_train_idx: ~/.mxnet/datasets/imagenet/rec/train.idx
  rec_val: ~/.mxnet/datasets/imagenet/rec/val.rec
  rec_val_idx: ~/.mxnet/datasets/imagenet/rec/val.idx
  resume_epoch: 0
  start_epoch: 0
  teacher: null
  temperature: 20
  transfer_lr_mult: 0.01
  use_rec: false
  warmup_epochs: 0
  warmup_lr: 0.0
  wd: 0.0001
valid:
  batch_size: 128
  num_workers: 4

Let’s modify the network from resnet50_v1 to resnet18_v1b

import fileinput
with fileinput.FileInput(config_name,
                         inplace = True, backup ='.bak') as f:
    for line in f:
        if 'resnet50_v1' in line:
            new_line = line.replace('resnet50_v1', 'resnet18_v1b')
            print(new_line, end='')
        else:
            print(line, end='')

The new classifier now should reflect the new configs we just edited

new_classifier2 = ImageClassificationEstimator(config_name)