pytorch-cifar
pytorch-cifar copied to clipboard
how to train my own dataset and classes?
-
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.
-
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.