Not able to retrain on mobilenet_v2 network

[jasvinderkhurana]

Hello,

I am trying to retrain classification model (mobilenet_v1 or mobilenet_v2), and I am following instructions at :

https://github.com/dusty-nv/jetson-inference/blob/master/docs/pytorch-cat-dog.md

In Help of train.py it shows architecture as below

  --model-dir MODEL_DIR
                        path to desired output directory for saving model
                        checkpoints (default: current directory)
  -a ARCH, --arch ARCH  model architecture: alexnet | densenet121 |
                        densenet161 | densenet169 | densenet201 | googlenet |
                        inception_v3 | mnasnet0_5 | mnasnet0_75 | mnasnet1_0 |
                        mnasnet1_3 | mobilenet_v2 | resnet101 | resnet152 |
                        resnet18 | resnet34 | resnet50 | resnext101_32x8d |
                        resnext50_32x4d | shufflenet_v2_x0_5 |
                        shufflenet_v2_x1_0 | shufflenet_v2_x1_5 |
                        shufflenet_v2_x2_0 | squeezenet1_0 | squeezenet1_1 |
                        vgg11 | vgg11_bn | vgg13 | vgg13_bn | vgg16 | vgg16_bn
                        | vgg19 | vgg19_bn | wide_resnet101_2 |
                        wide_resnet50_2 (default: resnet18)

We can see tha mobilnet_v2 is shown as supported. But when I execute below command:

python3 train.py -a mobilenet_v2 --model-dir=models/cat_dog data/cat_dog

I get Below error:

Use GPU: 0 for training
=> dataset classes:  2 ['cat', 'dog']
=> using pre-trained model 'mobilenet_v2'
Traceback (most recent call last):
  File "train.py", line 506, in <module>
    main()
  File "train.py", line 135, in main
    main_worker(args.gpu, ngpus_per_node, args)
  File "train.py", line 205, in main_worker
    model = reshape_model(model, args.arch, num_classes)
  File "/jetson-inference/python/training/classification/reshape.py", line 55, in reshape_model
    print("classifier reshaping not supported for " + args.arch)
NameError: name 'args' is not defined

Can anybody please guide How can I retrain mobilenet classification network for Nvidia Jetson devices?