Hello,
I am new to PyTorch, I just want to ensure that I correctly understand how model.to(device=device)
works before I start creating more complex models.
So far, almost all the tutorials that Ive followed create, train and evaluate their models in the same notebook/script.
However, what if I am creating my model, training it, and evaluating it using functions/operations in different scripts?
Suppose I have the following code under these different scripts.
main.py
model = EfficientNet.from_pretrained("efficientnet-b3")
model = model.to(device=device)
for epoch in range(EPOCHS):
train(model, train_loader, ...)
check_accuracy(model, val_loader, ...)
train_model.py
def train_model(model, loader, ...):
losses = []
for batch_idx, (data, targets) in enumerate(tqdm(loader)):
...
scores=model(data)
...
print(f"Loss: {sum(losses)/len(losses)}")
evaluation.py
def check_accuracy(model, loader, ...):
model.eval()
with torch.no_grad():
...
...
model.train()
return accuracy
Since my model is moved to device
by running .to(device=device)
, I dont need to return the model from each function in different modules to main.py
, is that right?
Simply passing the model to train_model()
and check_accuracy()
is enough run operations on the model, and once they finish running, the remaining operations in main.py will resume working on the same model, is that right?
Thank you all for your help.