File size: 1,835 Bytes
97daae4 cf617f9 97daae4 |
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 |
import torch
import torch.nn as nn
from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights
from torchmetrics.classification import MulticlassAccuracy
import data, engine, utils
# if __name__ == "__main__":
# HYPERPARAMETERS
SEED = 64
NUM_EPOCH = 50
LEARNIGN_RATE = 4e-6 # 3e-4, 4e-5, 7e-6, 5e-7, 3e-9
NUM_CLASSES = 101
# CUDA_LAUNCH_BLOCKING=1
device = torch.device("cuda:3" if torch.cuda.is_available() else 'cpu')
# print(device)
if __name__ == "__main__":
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
model = efficientnet_b0(weights=EfficientNet_B0_Weights.DEFAULT)
model.classifier = nn.Sequential(
nn.Dropout(p = 0.2, inplace = True),
nn.Linear(1280, NUM_CLASSES),
# nn.Softmax()
)
model = model.to(device)
# print(model)
loss_fn = torch.nn.CrossEntropyLoss(label_smoothing=0.1)
accuracy_fn = MulticlassAccuracy(num_classes = NUM_CLASSES).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = LEARNIGN_RATE)
# print(f"lr: {scheduler.optimizer.param_groups[0]['lr']}")
train_losses, test_losses, train_accs, test_accs, train_model = engine.train(model, data.train_dataloader, data.test_dataloader,
optimizer, loss_fn, accuracy_fn, NUM_EPOCH, device)
utils.save_model(model = train_model, target_dir = "./save_model", model_name = f"train_model_{LEARNIGN_RATE}.pth")
utils.plot_graph(train_losses = train_losses, test_losses = test_losses, train_accs = train_accs,
test_accs = test_accs, fig_name = f"plots/cnn_train_Loss_and_accuracy_plot_{LEARNIGN_RATE}.jpg") |