View file src/colab/nnufa_pytorch_new.py - Download
# -*- coding: utf-8 -*-
"""nnufa_pytorch_new.ipynb
Automatically generated by Colaboratory.
Original file is located at
https://colab.research.google.com/drive/1igFiqS2qsVlKFVHe26I6prtfZVqMQESm
Neural networks are universal functions approximators
"""
# Commented out IPython magic to ensure Python compatibility.
import torch
from torch import nn
import numpy
# %matplotlib inline
import matplotlib.pyplot as plt
plt.style.use('seaborn-whitegrid')
import numpy as np
# Get cpu or gpu device for training.
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
def f(x):
return 0.5 * x ** 3 + x ** 2 - 5 * x + 6
fig = plt.figure()
ax = plt.axes()
x = np.linspace(-4, 4, 1000)
ax.plot(x, f(x))
# Define model
N = 10
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
self.linear_relu_stack = nn.Sequential(
nn.Linear(1, N),
nn.ReLU(),
# nn.Linear(N, N),
# nn.ReLU(),
nn.Linear(N, 1),
)
def forward(self, x):
return self.linear_relu_stack(x)
model = NeuralNetwork().to(device)
print(model)
print(list(model.parameters()))
def fmodel(x):
x = float(x)
x = [[x]]
x = torch.tensor(x).to(device)
y = model(x)
# return y
return float(y[0][0])
for x in range(-4, 4):
y = fmodel(x)
print(f"{x} : {y}")
xs = np.linspace(-4, 4, 81)
print(f"xs = {xs}")
inputs = torch.tensor([[x] for x in xs], dtype=torch.float).to(device)
print(f"inputs:{inputs.shape}")
target = [[f(x)] for x in xs]
print(f"target = {target}")
target_tensor = torch.tensor(target, requires_grad=True).to(device)
print(f"target_tensor:{target_tensor.shape}")
def test():
with torch.no_grad():
ys = [fmodel(x) for x in xs]
# print(ys)
plt.plot(xs, ys, color='red')
target_plot = [f(x) for x in xs]
plt.plot(xs, target_plot, color='green')
test()
def compute_loss(prediction):
loss_tensor = prediction - target_tensor
loss_tensor = torch.flatten(loss_tensor)
loss = torch.dot(loss_tensor, loss_tensor)
return loss
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
def train():
model.train()
prediction = model(inputs)
loss = compute_loss(prediction)
# loss = loss_fn(prediction, target_tensor)
optimizer.zero_grad()
loss.backward()
optimizer.step()
return loss
for epoch in range(10000):
loss = train()
if epoch % 1000 == 0:
print(f"{epoch}: loss={loss}")
test()