使用 Optimizer

regression_torch3.py

import torch

import cv2
import matplotlib.pyplot as plt

def model(x, w, b):
    return w * x + b
    
def mse_loss(p, y):
    return ((p - y) ** 2).mean()
    
def training_loop(epochs, lr, params, x, y, verbose = False):
    optimizer = torch.optim.SGD([params], lr = lr)
    
    for epoch in range(1, epochs + 1):
        optimizer.zero_grad()
             
        w, b = params
        p = model(x, w, b)
        loss = mse_loss(p, y)
        loss.backward()
        optimizer.step()
           
        if verbose:
            print('週期', epoch, '--')
            print('\t損失：', float(loss))
            print('\t模型參數：', params)
            
    return params.detach()
  
# https://openhome.cc/Gossip/DCHardWay/images/PolynomialRegression-1.JPG
img = torch.from_numpy(cv2.imread('PolynomialRegression-1.JPG', cv2.IMREAD_GRAYSCALE))
idx = torch.where(img < 127)       # 黑點的索引

x = idx[1]
y = -idx[0] + img.shape[0] # 反轉 y 軸

plt.gca().set_aspect(1)
plt.scatter(x, y)

w, b = training_loop(
    epochs = 100, 
    lr = 0.001, 
    params = torch.tensor([1.0, 0.0], requires_grad = True), 
    x = x, 
    y = y
)

x = torch.linspace(0, 50, 50)
y = w * x + b

plt.plot(x, y)
plt.show()