# Estrutura básica para projetos de Machine Learning e Deep Learning
# Por Adriano Santos.

from torch import nn, relu
import torch.nn.functional as F
import torch.optim as optim
import torch
from torchvision import models


class ResNet(nn.Module):
    def __init__(self, saida, pretreinado=True):
        super(ResNet, self).__init__()

        resnet = models.resnet34(pretrained=pretreinado)
        layers = list(resnet.children())[:8]

        self.features1 = nn.Sequential(*layers[:6])
        self.features2 = nn.Sequential(*layers[6:])

        self.classificador = nn.Sequential(nn.BatchNorm1d(512), nn.Linear(512, saida))

    def forward(self, x):
        x = self.features1(x)
        x = self.features2(x)
        x = F.relu(x)
        x = nn.AdaptiveAvgPool2d((1,1))(x)
        x = x.view(x.shape[0], -1)
        return self.classificador(x)
[((15, 17, 15, 56), 'torchvision.models.resnet34', 'models.resnet34', (), '', False, 'from torchvision import models\n'), ((18, 25, 18, 51), 'torch.nn.Sequential', 'nn.Sequential', ({(18, 39, 18, 50): '*layers[:6]'}, {}), '(*layers[:6])', False, 'from torch import nn, relu\n'), ((19, 25, 19, 51), 'torch.nn.Sequential', 'nn.Sequential', ({(19, 39, 19, 50): '*layers[6:]'}, {}), '(*layers[6:])', False, 'from torch import nn, relu\n'), ((26, 12, 26, 21), 'torch.nn.functional.relu', 'F.relu', ({(26, 19, 26, 20): 'x'}, {}), '(x)', True, 'import torch.nn.functional as F\n'), ((21, 43, 21, 62), 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', ({(21, 58, 21, 61): '512'}, {}), '(512)', False, 'from torch import nn, relu\n'), ((21, 64, 21, 85), 'torch.nn.Linear', 'nn.Linear', ({(21, 74, 21, 77): '512', (21, 79, 21, 84): 'saida'}, {}), '(512, saida)', False, 'from torch import nn, relu\n'), ((27, 12, 27, 39), 'torch.nn.AdaptiveAvgPool2d', 'nn.AdaptiveAvgPool2d', ({(27, 33, 27, 38): '(1, 1)'}, {}), '((1, 1))', False, 'from torch import nn, relu\n')]