lenet5 / utils.py
HudsonArauj
updade
7557669
import os
import torch
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.transforms.functional import to_pil_image
import PIL
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.metrics import accuracy_score
arquivo = 'lenet-modular2024-06-26.pt'
class Lenet5Modular(nn.Module):
def __init__(self):
super(Lenet5Modular,self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3,out_channels=6,kernel_size=(5,5),stride=(1,1)),
nn.Tanh(),
nn.AvgPool2d(kernel_size=(2,2),stride=(2,2))
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=6,out_channels=16,kernel_size=(5,5),stride=(1,1)),
nn.Tanh(),
nn.AvgPool2d(kernel_size=(2,2),stride=(2,2))
)
self.conv3 = nn.Sequential(
nn.Conv2d(in_channels=16,out_channels=120,kernel_size=(5,5),stride=(1,1)),
nn.Tanh()
)
self.flat = nn.Flatten()
self.linear1 = nn.Sequential(
nn.Linear(120,84),
nn.Tanh()
)
self.linear2 = nn.Linear(84,10)
def forward(self,x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = self.flat(x)
x = self.linear1(x)
x = self.linear2(x)
return x
modelo = Lenet5Modular()
# modelo.load_state_dict(torch.load(arquivo))
modelo.cpu()
a = torch.rand(1, 3, 32, 32)
b = modelo(a)
b.size()
# Função para criar o transformador
def create_transformer():
return transforms.Compose([
transforms.Resize((32, 32)), # Redimensionar para 32x32
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(10),
transforms.RandomAffine(0, shear=10, scale=(0.8,1.2)),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
transforms.ToTensor()
])