54 lines
1.3 KiB
Python
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)
|