pytorch-cifar icon indicating copy to clipboard operation
pytorch-cifar copied to clipboard

how to train my own dataset and classes?

Open henbucuoshanghai opened this issue 2 years ago • 1 comments

henbucuoshanghai avatar Apr 18 '22 08:04 henbucuoshanghai

  1. In main.py right after transforms and before actually creating a model, replace that code with your own dataset class and its data-loader, followed by classes in your dataset.

  2. you'll have to modify the models to take in an extra parameter that defines your number of classes and replace the '10' in the last linear layer of models with that parameter. voila.

logan-mo avatar Aug 06 '22 17:08 logan-mo