AI_PRO/neuralNetwork.py
2021-06-22 22:56:01 +02:00

54 lines
1.3 KiB
Python

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from matplotlib.pyplot import imshow
import numpy as np
from matplotlib.pyplot import imshow
import matplotlib.pyplot as ppl
def plotdigit(image):
img = np.reshape(image, (-250, 250))
imshow(img, cmap='Greys', vmin=0, vmax=255)
ppl.show()
train_images = np.load('neural_network\\image-dataset.npy')
print(train_images.shape)
train_labels = np.load('neural_network\\label-dataset.npy')
train_images = train_images / 255
train_labels = train_labels / 255
train_images = [torch.tensor(image, dtype=torch.float32) for image in train_images]
print(train_images[0].shape)
train_labels = [torch.tensor(label, dtype=torch.long) for label in train_labels]
input_dim = 100*100*3
output_dim = 2
model = nn.Sequential(
nn.Linear(input_dim, output_dim),
nn.LogSoftmax()
)
def train(model, n_iter):
criterion = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
for epoch in range(n_iter):
for image, label in zip(train_images, train_labels):
optimizer.zero_grad()
output = model(image)
loss = criterion(output.unsqueeze(0), label.unsqueeze(0))
loss.backward()
optimizer.step()
print(f'epoch: {epoch:03}')
train(model, 100)