Custom dataset training for mobilenet-ss?

[donipologhimire]

Env

• GPU, e.g. NVIDIA Xavier AGX
• OS, e.g. Ubuntu18.04
• Cuda version : 10.2.89
• TensorRT version

About this repo

• which branch/tag/commit are you using? https://github.com/dusty-nv/pytorch-detection
• which model? mobile net ssd

Your problem

I wanted to custom train on a different dataset (eg, doorways detection).

• I see that the train_ssd.py seems to do it, however I have not not been able to realize how to fo it.
• I am new to computer vision, can you guide me how the dataset should look like, I have images (.jpeg) and label (.xml files).
• The folder is structures as \images - > \train and \test . these subfolders have images and labels.
• Do I need to load a model in order to do it, because I see that is the case, however I donot have any model.
• How do I prune my model after I have created one?
• How to convert the .pth file to trt model once I am able to train it.

Thank you.

[huzaifa-jawad367]

once you have your data prepared, resolved dependencies and models prepared you can train the modal by using the following command

python3 train_ssd.py --data=data/fruit --model-dir=models/fruit --batch-size=4 --epochs=30

Regarding your issue with the format of the dataset, the pytorch-ssd model supports 2 ways of reading data one is in OpenImages format which is default and the other format supported is the VOC format. XML files are used in VOC format you can specify the dataset type in the terminal command by adding --dataset-type=voc.
The structure of your directory should be as follows

- dataset_directory
  - annotations
  - JPEGImages
  - SegmentationClass (include if performing segmentation)
  - SegmentationObject (include if performing segmentation)
  - ImageSets
    - Main
      - test.txt
      - trainval.txt

you can include the path to the dataset by including the following in your terminal command --data=<path/to/your/dataset_directory>

pytorch-ssd supports the following models:
mb1-ssd, mb1-ssd-lite, mb2-ssd-lite or vgg16-ssd
you can specify what model you want by using the `--net=' in your training command.
you can find the pre-trained weights for one of the model here and find for other architectures online.

you can convert the model to onnx using the following command:

python3 onnx_export.py --model-dir=models/fruit

for further detail look at this