Clean the deep learning notebook code to .py files and for production

Question

I am using a code (notebook) to predict. Now I should clean the code and use it into production (change to .py file and import and etc). I am using a class for my model as follow:

class Autoencoder(nn.Module):

    def __init__(self, input_dim, par_dim):
        super().__init__()

    def encode(self,y):
    return x

    def forward():....
    def test():...

And the train batch and training is as follow (they are not in the class):

def train_batch(model, optimizer, device, batch, labels):

def train(model, device, epochs, train_iterator, optimizer, validate_iterator):

And then I consider the model as:

model              = Autoencoder()

optimizer          = torch.optim.Adam(model.parameters() )

loss = train(model, device, epochs, train_iterator, optimizer, validate_iterator) 

So, my question is: I should use one .py file for the class? and one .py file for the train, and one .py file for the train_batch?

Or I can use this train and train_batch as inside the class?

Do you have any tutorial (video or stackoverflow link) for this type of the work? Thank you. Thank you

Answer 1

You can use 1 file only and import the classes or functions when needed.

let's say they are all in the file foo.py

you can do:

from foo import Autoencoder, train_batch, train

In another file given they are in the same folder or the foo.py is in the pythonpath