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()