-
Notifications
You must be signed in to change notification settings - Fork 0
/
All_Train_A05.py
79 lines (62 loc) · 2.71 KB
/
All_Train_A05.py
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
###############################################################################
# IMPORTS
###############################################################################
import os
import sys
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import v2
from sklearn.metrics import (accuracy_score, f1_score)
import time
import A05
base_dir = "all_assign05"
out_dir = base_dir + "/" + "output"
###############################################################################
# MAIN
###############################################################################
def main():
# Create output directory
os.makedirs(out_dir, exist_ok=True)
# Get names of all approaches
all_names = A05.get_approach_names()
chosen_approach_names = all_names
for approach_name in chosen_approach_names:
print("TRAINING APPROACH:", approach_name)
# Create data transforms
train_transform = A05.get_data_transform(approach_name, training=True)
test_transform = A05.get_data_transform(approach_name, training=False)
# Load CIFAR10 data
training_data = datasets.CIFAR10(root="data", train=True, download=True, transform=train_transform)
test_data = datasets.CIFAR10(root="data", train=False, download=True, transform=test_transform)
# Create dataloaders
batch_size = A05.get_batch_size(approach_name)
train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=batch_size)
# Set number of classes
class_cnt = 10
# Create the model
model = A05.create_model(approach_name, class_cnt)
print("MODEL:", approach_name)
print(model)
# Move to GPU if possible
device = ("cuda" if torch.cuda.is_available()
else "mps" if torch.backends.mps.is_available()
else "cpu"
)
print(f"Using {device} device")
model = model.to(device)
# Train classifiers
start_time = time.time()
print("Training " + approach_name + "...")
model = A05.train_model(approach_name, model, device, train_dataloader, test_dataloader)
print("Training complete!")
print("Time taken:", (time.time() - start_time))
# Save the model
model_path = os.path.join(out_dir, "model_" + approach_name + ".pth")
torch.save(model.state_dict(), model_path)
print("Model saved to:", model_path)
if __name__ == "__main__":
main()