my_model_fc

import torch
from torch import nn
import torch.nn.functional as F

class Network(nn.Module):
    def __init__(self, input_size, output_size, hidden_layers, drop_p = 0.2):
        super().__init__()

        self.hidden_layers = nn.ModuleList([nn.Linear(input_size, hidden_layers[0])])

        self.layers_size = zip(hidden_layers[:-1], hidden_layers[1:])

        self.hidden_layers.extend([nn.Linear(h1, h2) for h1, h2 in self.layers_size])

        self.output = nn.Linear(hidden_layers[-1], output_size)

        self.dropout = nn.Dropout(p = drop_p)


    def forward(self, x):
        x = x.view(x.shape[0], -1)

        for each in self.hidden_layers:
            x = each(x)
            x = F.relu(x)
            x = self.dropout(x)

        x = self.output(x)

        return F.log_softmax(x, dim=1)

def validation(model, testloader, criterion):

    accuracy = 0
    test_loss = 0

    for images, labels in testloader:
        log_ps = model.forward(images)
        test_loss += criterion(log_ps, labels)
        ps = torch.exp(log_ps)

        top_ps, top_class = ps.topk(1, dim=1)

        equals  = top_class == labels.view(*top_class.shape)

        accuracy += torch.mean(equals.type(torch.FloatTensor))

    return test_loss, accuracy


def train(model, trainloader, testloader, criterion, optimizer, epochs=2):

    train_losses, test_losses = [], []
    for e in range (epochs):
        train_loss = 0
        for images, labels in trainloader:

            optimizer.zero_grad()

            log_ps = model.forward(images)
            loss = criterion(log_ps, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
        else:
            with torch.no_grad():
                model.eval()

                test_loss, accuracy = validation(model, testloader, criterion)

            model.train()

            train_losses.append(train_loss/len(trainloader))
            test_losses.append(test_loss/len(testloader))

            print(f'Epoch: {e+1}/{epochs}')
            print('Training Loss: {:.3f}'.format(train_loss/len(trainloader)))
            print('Test Loss: {:.3f}'.format(test_loss/len(testloader)))
            print('Accuracy: {:.3f}'.format(accuracy/len(testloader)))
            print()