File size: 3,087 Bytes
7fc0372 9a437fa 7fc0372 9a437fa 7fc0372 9a437fa 7fc0372 9a437fa 7fc0372 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
"""
Trains a Pytorch image classification model using device agonstic code
"""
import os
import torch
import data_setup, engine, model_builder, utils
from torchvision import transforms
from timeit import default_timer as timer
from pathlib import Path
from get_data import fetch_data
import json
def train_torch():
# Load config.json
with open("config.json", "r") as f:
config = json.load(f)
# Access configuration parameters
input_shape = config["input_shape"]
hidden_units = config["hidden_units"]
# Setup hyperparameters
NUM_EPOCHS = config["num_epochs"]
BATCH_SIZE = config["batch_size"]
LEARNING_RATE = config["learning_rate"]
# Define the URL and paths
# Define the URL and paths
url = "https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip"
data_path = Path("/data")
image_path = Path(data_path / "pizza_steak_sushi")
# Call the fetch_data function
fetch_data(url, data_path, image_path)
# Setup directories
train_dir = "/data/pizza_steak_sushi/train"
test_dir = "/data/pizza_steak_sushi/test"
# Setup target device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Create transforms
data_transform = transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor()
])
# Create dataloaders with help from data_Setup.py
train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(train_dir=train_dir,
test_dir=test_dir,
transform=data_transform,
batch_size=BATCH_SIZE)
# Create model with help from model_builder.py
model = model_builder.TinyVGG(
input_shape = input_shape,
hidden_units=hidden_units,
output_shape=len(class_names)
).to(device)
# Set loss and optimizer
loss_fn = torch.nn.CrossEntropyLoss()
optimizer= torch.optim.Adam(model.parameters(),
lr=LEARNING_RATE)
# Start the timer
start_time = timer()
# Start training with help from engine.py
results = engine.train(model=model,
train_dataloader=train_dataloader,
test_dataloader=test_dataloader,
loss_fn=loss_fn,
optimizer=optimizer,
epochs=NUM_EPOCHS,
device=device)
# End the timer and print out how long it took
end_time = timer()
print(f"[INFO] Total training time: {end_time-start_time:.3f} seconds")
print(f"[INFO] results dict: {results}")
# Save the model with help from utils.py
utils.save_model(model=model,
target_dir="models",
model_name="05_going_modular_script_mode_tinyvgg_model.pth")
if __name__ == '__main__':
train_torch()
|