JETSON NANO - CREA TU PRIMER RED NEURONAL CONVOLUCIONAL (CNN)

[@IN3DMEX2.0]

#MODELO import torch import torch.nn as nn class CNN(nn.Module): def __init__(self): super(CNN, self).__init__() self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1) self.relu1 = nn.ReLU() self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2) self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1) self.relu2 = nn.ReLU() self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2) self.flatten = nn.Flatten() self.fc1 = nn.Linear(32 * 56 * 56, 128) self.relu3 = nn.ReLU() self.fc2 = nn.Linear(128, 2) def forward(self, x): x = self.conv1(x) x = self.relu1(x) x = self.maxpool1(x) x = self.conv2(x) x = self.relu2(x) x = self.maxpool2(x) x = self.flatten(x) x = self.fc1(x) x = self.relu3(x) x = self.fc2(x) return x # Crea una instancia de la CNN modelo = CNN()

[@IN3DMEX2.0]

#ENTRENAMIENTO import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms from modelo import CNN # Importa la clase CNN # Define el dispositivo (CPU o GPU) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Define las transformaciones para las imágenes data_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # Carga el dataset de entrenamiento train_dataset = datasets.ImageFolder(root='data/training_set', transform=data_transform) # Crea un DataLoader para el dataset de entrenamiento train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) # Crea una instancia de la CNN y mueve la red al dispositivo modelo = CNN().to(device) # Define el optimizador y la función de pérdida optimizer = optim.Adam(modelo.parameters(), lr=0.001) criterion = nn.CrossEntropyLoss() # Entrena la red num_epochs = 10 for epoch in range(num_epochs): for i, (images, labels) in enumerate(train_loader): images = images.to(device) labels = labels.to(device) outputs = modelo(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() if i % 100 == 0: print(f'Epoch {epoch+1}/{num_epochs}, Batch {i+1}/{len(train_loader)}, Loss: {loss.item():.4f}') # Guarda los pesos del modelo entrenado torch.save(modelo.state_dict(), 'modelo_pesos.pth')

[@IN3DMEX2.0]

#INFERENCIA import torch from torchvision import datasets, transforms from modelo import CNN # Importa la clase CNN import os from PIL import Image # Define el dispositivo (CPU o GPU) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Define las transformaciones para las imágenes data_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # Crea una instancia de la CNN y mueve la red al dispositivo modelo = CNN().to(device) # Carga los pesos del modelo entrenado modelo.load_state_dict(torch.load('modelo_pesos.pth')) # Pon el modelo en modo de evaluación modelo.eval() # Realiza predicciones sobre nuevas imágenes with torch.no_grad(): image_path = os.path.abspath('data/test_set/cats/cat.4001.jpg') image = Image.open(image_path) # Aplica las transformaciones a la imagen image = data_transform(image) # Convierte la imagen a un tensor y mueve al dispositivo image = image.unsqueeze(0).to(device) # Realiza la predicción output = modelo(image) # Obtiene la clase predicha _, predicted_class = torch.max(output.data, 1) # Define los nombres de las clases class_names = ['gato', 'perro'] # Obtiene el nombre de la clase predicted_class_name = class_names[predicted_class.item()] # Imprime la clase predicha print(f'Predicted class: {predicted_class_name}')