First nn with simple linear network

.gitignore

+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# IPython Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# dotenv
+.env
+
+# virtualenv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+
+# Rope project settings
+.ropeproject
+### VirtualEnv template
+# Virtualenv
+# http://iamzed.com/2009/05/07/a-primer-on-virtualenv/
+.Python
+[Bb]in
+[Ii]nclude
+[Ll]ib
+[Ll]ib64
+[Ll]ocal
+[Ss]cripts
+pyvenv.cfg
+.venv
+pip-selfcheck.json
+### JetBrains template
+# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and Webstorm
+# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
+
+# User-specific stuff:
+.idea/workspace.xml
+.idea/tasks.xml
+.idea/dictionaries
+.idea/vcs.xml
+.idea/jsLibraryMappings.xml
+
+# Sensitive or high-churn files:
+.idea/dataSources.ids
+.idea/dataSources.xml
+.idea/dataSources.local.xml
+.idea/sqlDataSources.xml
+.idea/dynamic.xml
+.idea/uiDesigner.xml
+
+# Gradle:
+.idea/gradle.xml
+.idea/libraries
+
+# Mongo Explorer plugin:
+.idea/mongoSettings.xml
+
+.idea/
+
+## File-based project format:
+*.iws
+
+## Plugin-specific files:
+
+# IntelliJ
+/out/
+
+# mpeltonen/sbt-idea plugin
+.idea_modules/
+
+# JIRA plugin
+atlassian-ide-plugin.xml
+
+# Crashlytics plugin (for Android Studio and IntelliJ)
+com_crashlytics_export_strings.xml
+crashlytics.properties
+crashlytics-build.properties
+fabric.properties
+
+# Dataset
+data/

main.py

+from os.path import isfile
+
+import torch
+from numpy import prod
+from torch import nn
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from torchvision.transforms import ToTensor
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using {device} device")
+
+
+def get_data(batch_size: int = 64):
+    # Download training data from open datasets.
+    training_data = datasets.CIFAR10(
+        root="/home/flifloo/IA/data",
+        train=True,
+        download=True,
+        transform=ToTensor(),
+    )
+
+    # Download test data from open datasets.
+    testing_data = datasets.CIFAR10(
+        root="/home/flifloo/IA/data",
+        train=False,
+        download=True,
+        transform=ToTensor(),
+    )
+
+    # Create data loaders.
+    train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
+    test_dataloader = DataLoader(testing_data, batch_size=batch_size, shuffle=True)
+
+    return train_dataloader, test_dataloader
+
+
+def generate_layers(inp: int, output: int):
+    layers = 2
+    conns = (inp+output)*2
+    stack = [nn.Linear(inp, conns), nn.ReLU()]
+
+    print(f"input: {inp}, output: {output}, layers: {layers}, conns: {conns}")
+
+    print("Generating stack...")
+    for _ in range(layers):
+        stack.append(nn.Linear(conns, conns))
+        stack.append(nn.ReLU())
+
+    stack += [nn.Linear(conns, output), nn.ReLU()]
+
+    print("Stack generated")
+    return stack
+
+
+# Define model
+class NeuralNetwork(nn.Module):
+    def __init__(self, stack):
+        super(NeuralNetwork, self).__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(*stack)
+
+    def forward(self, x):
+        return self.linear_relu_stack(self.flatten(x))
+
+
+def train(dataloader, model, loss_fn, optimizer):
+    size = len(dataloader.dataset)
+    for batch, (X, y) in enumerate(dataloader):
+        X, y = X.to(device), y.to(device)
+
+        # Compute prediction error
+        pred = model(X)
+        loss = loss_fn(pred, y)
+
+        # Backpropagation
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        if batch % 100 == 0:
+            loss, current = loss.item(), batch * len(X)
+            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
+
+
+def test(dataloader, model, loss_fn):
+    size = len(dataloader.dataset)
+    model.eval()
+    test_loss, correct = 0, 0
+    with torch.no_grad():
+        for X, y in dataloader:
+            X, y = X.to(device), y.to(device)
+            pred = model(X)
+            test_loss += loss_fn(pred, y).item()
+            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
+    test_loss /= size
+    correct /= size
+    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
+    return correct
+
+
+def training():
+    train_data, test_data = get_data()
+
+    stack = generate_layers(prod(test_data.dataset.data[0].shape), len(test_data.dataset.classes))
+    model = NeuralNetwork(stack).to(device)
+    if isfile("model.pth"):
+        print("Loading model from save")
+        model.load_state_dict(torch.load("model.pth"))
+
+    print(model)
+
+    loss_fn = nn.CrossEntropyLoss()
+    # lr = sur/sous appretisage
+    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=.9)
+
+    e = 0
+    c = 0
+    while c < 0.90:
+        print(f"Epoch {e+1}\n-------------------------------")
+        train(train_data, model, loss_fn, optimizer)
+        c = test(test_data, model, loss_fn)
+        torch.save(model.state_dict(), "model.pth")
+        e += 1
+    print("Done!")
+
+
+if __name__ == '__main__':
+    training()