Refactor training loop from script to class #40
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 is currently in the
main.py
file, is encapsulated in a new class namedMNISTTrainer
. The code for loading and preprocessing the MNIST dataset is moved into a method namedload_data
, and the code for defining the PyTorch model is moved into a method nameddefine_model
. Additionally, a new method namedtrain
is added toMNISTTrainer
to contain the training loop.Summary of Changes
MNISTTrainer
inmain.py
.load_data
inMNISTTrainer
.define_model
inMNISTTrainer
.train
inMNISTTrainer
to contain the training loop.api.py
to importMNISTTrainer
instead ofNet
.api.py
to use theload_data
anddefine_model
methods ofMNISTTrainer
.Fixes #6.
🎉 Latest improvements to Sweep:
💡 To get Sweep to edit this pull request, you can: