Programmiererforum

Code: Select all

import torch
import torch.nn as nn
import torch.optim as optim
from tqdm.auto import tqdm
from torchvision import models

# Set device (GPU if available)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Define the ResNet18 model
class ResNet18(nn.Module):
def __init__(self, num_classes):
super(ResNet18, self).__init__()
self.model = models.resnet18(pretrained=False)  # Set to True to use pre-trained weights
self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)  # Modify the final layer

def forward(self, x):
return self.model(x)

# Initialize the model with the number of classes in your dataset
num_classes = len(train_data_simple.classes)  # Adjust this based on your augmented dataset
model_resnet = ResNet18(num_classes=num_classes).to(device)

# Loss function and optimizer (SGD and CrossEntropy)
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_resnet.parameters(), lr=0.1, momentum=0.9)

# Training loop (with progress bar using tqdm)
def train(model, train_dataloader, test_dataloader, optimizer, loss_fn, epochs):
for epoch in tqdm(range(epochs)):
model.train()  # Set the model to training mode
running_loss = 0.0
correct_train, total_train = 0, 0

for batch, (X, y) in enumerate(train_dataloader):
X, y = X.to(device), y.to(device)

optimizer.zero_grad()  # Zero the gradients

outputs = model(X)  # Forward pass
loss = loss_fn(outputs, y)  # Compute loss

loss.backward()  # Backpropagation
optimizer.step()  # Update the weights

running_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total_train += y.size(0)
correct_train += (predicted == y).sum().item()

# Print training loss and accuracy
print(f"Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(train_dataloader):.4f}")
print(f"Training Accuracy: {100 * correct_train / total_train:.2f}%")

# Validation
model.eval()  # Switch model to evaluation mode for validation
correct_val, total_val = 0, 0
with torch.no_grad():
for X_val, y_val in test_dataloader:
X_val, y_val = X_val.to(device), y_val.to(device)
val_outputs = model(X_val)
_, predicted = torch.max(val_outputs.data, 1)
total_val += y_val.size(0)
correct_val += (predicted == y_val).sum().item()

# Print validation accuracy
print(f"Validation Accuracy: {100 * correct_val / total_val:.2f}%")

return model

# Example training
NUM_EPOCHS = 30  # Adjust epochs as per your requirement
# Use the augmented training data loader
trained_model_resnet = train(model_resnet, train_dataloader_simple, test_dataloader_simple, optimizer, loss_fn, NUM_EPOCHS)

< /code>
Dies ist der Code, den ich in einer verschiedenen Zelle ausgeführt habe, um das Modell < /p>
zu speichern# Save the trained model
model_path = "trained_resnet18_model.pth"
torch.save(model_resnet.state_dict(), model_path)
print(f"Model saved to {model_path}")