1. Configuração do Disopsitivo

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision import datasets
import pathlib, warnings

warnings.filterwarnings("ignore")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

2. Carregamento dos Dados

data_dir = pathlib.Path('./data/8-data')
class_names = [p.name for p in data_dir.iterdir() if p.is_dir()]

transform = transforms.Compose([
    transforms.Resize([224, 224]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

dataset = datasets.ImageFolder('./data/8-data/', transform=transform)

3. Divisão do Conjunto de Dados

train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_set, test_set = torch.utils.data.random_split(dataset, [train_size, test_size])

batch_size = 4
train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=3)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=3)

Componentes Básicos

def autopad(kernel, padding=None):
    if padding is None:
        padding = kernel // 2 if isinstance(kernel, int) else [k // 2 for k in kernel]
    return padding

class ConvBlock(nn.Module):
    def __init__(self, in_ch, out_ch, kernel=1, stride=1, padding=None, groups=1, activation=True):
        super().__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, kernel, stride, autopad(kernel, padding), groups=groups, bias=False)
        self.norm = nn.BatchNorm2d(out_ch)
        self.act = nn.SiLU() if activation else nn.Identity()

    def forward(self, x):
        return self.act(self.norm(self.conv(x)))

class BottleneckUnit(nn.Module):
    def __init__(self, in_ch, out_ch, use_shortcut=True, groups=1, expansion=0.5):
        super().__init__()
        hidden = int(out_ch * expansion)
        self.conv1 = ConvBlock(in_ch, hidden, 1, 1)
        self.conv2 = ConvBlock(hidden, out_ch, 3, 1, groups=groups)
        self.use_add = use_shortcut and in_ch == out_ch

    def forward(self, x):
        out = self.conv2(self.conv1(x))
        return x + out if self.use_add else out

class CSPBlock(nn.Module):
    def __init__(self, in_ch, out_ch, repeats=1, shortcut=True, groups=1, expansion=0.5):
        super().__init__()
        hidden = int(out_ch * expansion)
        self.branch1 = ConvBlock(in_ch, hidden, 1, 1)
        self.branch2 = ConvBlock(in_ch, hidden, 1, 1)
        self.merge = ConvBlock(2 * hidden, out_ch, 1, 1)
        self.blocks = nn.Sequential(*[BottleneckUnit(hidden, hidden, shortcut, groups, 1.0) for _ in range(repeats)])

    def forward(self, x):
        return self.merge(torch.cat([self.blocks(self.branch1(x)), self.branch2(x)], dim=1))

class FastSPP(nn.Module):
    def __init__(self, in_ch, out_ch, pool_kernel=5):
        super().__init__()
        hidden = in_ch // 2
        self.pre_conv = ConvBlock(in_ch, hidden, 1, 1)
        self.post_conv = ConvBlock(hidden * 4, out_ch, 1, 1)
        self.pool = nn.MaxPool2d(pool_kernel, stride=1, padding=pool_kernel // 2)

    def forward(self, x):
        x = self.pre_conv(x)
        p1 = self.pool(x)
        p2 = self.pool(p1)
        p3 = self.pool(p2)
        return self.post_conv(torch.cat([x, p1, p2, p3], dim=1))

Definição do Backbone do YOLOv5

class YOLOv5Backbone(nn.Module):
    def __init__(self, num_classes=4):
        super().__init__()
        self.stem = ConvBlock(3, 64, 3, 2, padding=2)
        self.stage1 = nn.Sequential(
            ConvBlock(64, 128, 3, 2),
            CSPBlock(128, 128)
        )
        self.stage2 = nn.Sequential(
            ConvBlock(128, 256, 3, 2),
            CSPBlock(256, 256)
        )
        self.stage3 = nn.Sequential(
            ConvBlock(256, 512, 3, 2),
            CSPBlock(512, 512)
        )
        self.stage4 = nn.Sequential(
            ConvBlock(512, 1024, 3, 2),
            CSPBlock(1024, 1024),
            FastSPP(1024, 1024)
        )
        self.classifier = nn.Sequential(
            nn.Linear(1024 * 8 * 8, 100),
            nn.ReLU(),
            nn.Linear(100, num_classes)
        )

    def forward(self, x):
        x = self.stem(x)
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)

model = YOLOv5Backbone().to(device)

Funções de Treinamento e Validação

def train_epoch(loader, model, criterion, optimizer):
    model.train()
    total_loss = 0
    correct = 0
    for inputs, targets in loader:
        inputs, targets = inputs.to(device), targets.to(device)
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        correct += (outputs.argmax(1) == targets).sum().item()
    acc = correct / len(loader.dataset)
    avg_loss = total_loss / len(loader)
    return acc, avg_loss

def validate(loader, model, criterion):
    model.eval()
    total_loss = 0
    correct = 0
    with torch.no_grad():
        for inputs, targets in loader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            total_loss += loss.item()
            correct += (outputs.argmax(1) == targets).sum().item()
    acc = correct / len(loader.dataset)
    avg_loss = total_loss / len(loader)
    return acc, avg_loss

Execução do Treinamento

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()
epochs = 20

train_acc_history, train_loss_history = [], []
val_acc_history, val_loss_history = [], []
best_accuracy = 0.0

for epoch in range(epochs):
    train_acc, train_loss = train_epoch(train_loader, model, criterion, optimizer)
    val_acc, val_loss = validate(test_loader, model, criterion)
    
    if val_acc > best_accuracy:
        best_accuracy = val_acc
        torch.save(model.state_dict(), 'best_model.pth')
    
    train_acc_history.append(train_acc)
    train_loss_history.append(train_loss)
    val_acc_history.append(val_acc)
    val_loss_history.append(val_loss)
    
    print(f'Epoch {epoch+1:2d} | Train Acc: {train_acc*100:.1f}% | Val Acc: {val_acc*100:.1f}%')

Gráficos de Desempenho

import matplotlib.pyplot as plt

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(train_acc_history, label='Treino')
plt.plot(val_acc_history, label='Validação')
plt.title('Acurácia por Época')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(train_loss_history, label='Treino')
plt.plot(val_loss_history, label='Validação')
plt.title('Perda por Época')
plt.legend()
plt.show()

Avaliação do Melhor Modelo

model.load_state_dict(torch.load('best_model.pth', map_location=device))
final_acc, final_loss = validate(test_loader, model, criterion)
print(f'Acurácia final: {final_acc:.4f}')