Refactor training loop from script to class #101
Closed
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR refactors the training loop from a script to a class for better organization and reusability. The training loop, which was previously located in
src/main.py
, has been encapsulated within a new class namedMNISTTrainer
. The class is responsible for loading and preprocessing the MNIST dataset, defining the PyTorch model, and training the model. It also provides a method to save the trained model to a file.Summary of Changes
MNISTTrainer
insrc/main.py
to encapsulate the training loop.load_data
inMNISTTrainer
.define_model
inMNISTTrainer
.train_model
inMNISTTrainer
to contain the training loop.save_model
inMNISTTrainer
to save the trained model to a file.src/api.py
to import theMNISTTrainer
class instead of theNet
class.src/api.py
to use theMNISTTrainer
class for loading the data, defining the model, training the model, and saving the model to a file.Please review and merge this PR. Thank you!
Fixes #6.
🎉 Latest improvements to Sweep:
💡 To get Sweep to edit this pull request, you can: