File size: 3,323 Bytes
99f802a ea56d2d 99f802a 7987245 99f802a ea56d2d 2e58968 b245442 ea56d2d 99f802a 7987245 99f802a ea56d2d 99f802a ea56d2d 99f802a ea56d2d 99f802a ea56d2d b245442 ea56d2d b245442 ea56d2d 2e58968 7987245 |
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 |
import os
import torch
import torch.optim as optim
import torch.nn as nn
from torchvision import datasets, models, transforms
from tqdm import tqdm
import torch
from consts import TRAIN_TEST_IMAGES_DIR
# Transformations for the image data
data_transforms = transforms.Compose([
transforms.Grayscale(num_output_channels=3), # Convert images to grayscale with 3 channels
transforms.RandomCrop((224, 224)), # Resize images to the expected input size of the model
transforms.ToTensor(), # Convert images to PyTorch tensors
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Normalize with ImageNet stats
])
# Create datasets
image_datasets = {
x: datasets.ImageFolder(os.path.join(TRAIN_TEST_IMAGES_DIR, x), data_transforms)
for x in ['train', 'test']
}
# Create dataloaders
dataloaders = {
'train': torch.utils.data.DataLoader(image_datasets['train'], batch_size=4, shuffle=True),
'test': torch.utils.data.DataLoader(image_datasets['test'], batch_size=4, shuffle=True)
}
# Define the model
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
# Modify the last fully connected layer to match the number of font classes you have
num_classes = len(image_datasets['train'].classes)
model.fc = nn.Linear(model.fc.in_features, num_classes)
# Define the loss function
criterion = torch.nn.CrossEntropyLoss()
# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
# Function to perform a training step with progress bar
def train_step(model, data_loader, criterion, optimizer):
model.train()
total_loss = 0
progress_bar = tqdm(data_loader, desc='Training', leave=True)
for inputs, targets in progress_bar:
outputs = model(inputs)
loss = criterion(outputs, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
progress_bar.set_postfix(loss=loss.item())
progress_bar.close()
return total_loss / len(data_loader)
# Function to perform a validation step with progress bar
def validate(model, data_loader, criterion):
model.eval()
total_loss = 0
correct = 0
progress_bar = tqdm(data_loader, desc='Validation', leave=False)
with torch.no_grad():
for inputs, targets in progress_bar:
outputs = model(inputs)
loss = criterion(outputs, targets)
total_loss += loss.item()
_, predicted = torch.max(outputs, 1)
correct += (predicted == targets).sum().item()
progress_bar.set_postfix(loss=loss.item())
progress_bar.close()
return total_loss / len(data_loader), correct / len(data_loader.dataset)
print(image_datasets['train'].classes)
# Training loop with progress bar for epochs
num_epochs = 10 # Replace with the number of epochs you'd like to train for
for epoch in range(num_epochs):
print(f"Epoch {epoch+1}/{num_epochs}")
train_loss = train_step(model, dataloaders["train"], criterion, optimizer)
val_loss, val_accuracy = validate(model, dataloaders["test"], criterion)
print(f"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val Accuracy: {val_accuracy:.4f}")
# Save the model to disk
torch.save(model.state_dict(), 'font_identifier_model.pth')
|