kuangliu / pytorch-cifar

95.47% on CIFAR10 with PyTorch
MIT License
5.94k stars 2.14k forks source link

how to train my own dataset and classes? #150

Closed henbucuoshanghai closed 4 months ago

logan-mo commented 2 years ago

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.