Template:NVIDIA Transfer Learning Toolkit - Training a Classifier

model_config {

  # Model architecture can be chosen from:
  # ['resnet', 'vgg', 'googlenet', 'alexnet', 'mobilenet_v1', 'mobilenet_v2', 'squeezenet']

  arch: "squeezenet"

  # for resnet --> n_layers can be [10, 18, 50]
  # for vgg --> n_layers can be [16, 19]

  # n_layers: 18 # Only relevant for resnet and vgg
  use_bias: True
  use_batch_norm: True
  all_projections: True
  use_pooling: False
  freeze_bn: False 
  freeze_blocks: 0 # When using pretrained-weights not all layers need to be retrained
  freeze_blocks: 1
  freeze_blocks: 2
  freeze_blocks: 3
  freeze_blocks: 4
  freeze_blocks: 5
  freeze_blocks: 6

  # image size should be "3, X, Y", where X,Y >= 16
  input_image_size: "3,112,112"
}

eval_config {
  eval_dataset_path: "test"
  model_path: "results/weights/squeezenet_080.tlt" # Has to be specific here
  top_k: 3 # If the correct class is in the top 3 (in this case) a special statistic is reported
  batch_size: 256
  n_workers: 8

}

train_config {
  train_dataset_path: "train"
  val_dataset_path: "val"
  pretrained_model_path: "tlt_squeezenet_classification_v1/squeezenet.hdf5"
  # optimizer can be chosen from ['adam', 'sgd']

  optimizer: "sgd"
  batch_size_per_gpu: 16
  n_epochs: 80
  n_workers: 16

  # regularizer
  reg_config {
    type: "L2"
    scope: "Conv2D,Dense"
    weight_decay: 0.00005

  }

  # learning_rate

  lr_config {

    # "step" and "soft_anneal" are supported.

    scheduler: "soft_anneal"

    # "soft_anneal" stands for soft annealing learning rate scheduler.
    # the following 4 parameters should be specified if "soft_anneal" is used.
    learning_rate: 0.005
    soft_start: 0.056
    annealing_points: "0.3, 0.6, 0.8"
    annealing_divider: 10
    # "step" stands for step learning rate scheduler.
    # the following 3 parameters should be specified if "step" is used.
    # learning_rate: 0.006
    # step_size: 10
    # gamma: 0.1
  }
}