66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
import torch
|
|
|
|
data_size = 20
|
|
|
|
|
|
torch.manual_seed(1234)
|
|
x_data = torch.abs(torch.randn((data_size)))
|
|
y_data = 10 * x_data
|
|
|
|
|
|
class MUL(torch.autograd.Function):
|
|
@staticmethod
|
|
def forward(ctx, input, weight):
|
|
ctx.save_for_backward(input, weight)
|
|
return input * weight
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
input, weight = ctx.saved_tensors
|
|
grad_weight = input * grad_output
|
|
grad_input = weight * grad_output
|
|
print(f"grad_output:{grad_output.item():.4f}")
|
|
return grad_input, grad_weight
|
|
|
|
|
|
class LinearModel(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = torch.nn.Parameter(torch.tensor([[1.0]]), requires_grad=True)
|
|
|
|
def forward(self, x):
|
|
return MUL.apply(x, self.weight)
|
|
return x * self.weight
|
|
|
|
|
|
model = LinearModel()
|
|
criterion = torch.nn.MSELoss()
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
|
|
loss_history = []
|
|
|
|
|
|
for step in range(data_size):
|
|
y_pred = model(x_data[step])
|
|
|
|
# loss = criterion(y_pred, y_data[step])
|
|
# loss = y_data[step] / y_pred - 1.0
|
|
# loss = torch.abs(y_data[step] - y_pred)
|
|
# loss = y_data[step] - y_pred
|
|
loss = (y_data[step] - y_pred) * (y_data[step] - y_pred)
|
|
|
|
loss_history.append(loss.item())
|
|
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
|
|
if (step + 1) % 1 == 0:
|
|
w = model.weight.item()
|
|
print(
|
|
f"Step {step+1}: w={w:.4f} loss={loss.item():.6f} input:{x_data[step]:.4f} output:{y_pred.item():.4f} label:{y_data[step].item():.4f} w_grad:{model.weight.grad.item():.4f}"
|
|
)
|
|
|
|
optimizer.step() # w = w - lr * ∇w[5](@ref)
|
|
|
|
test_x = torch.tensor([[5.0]])
|
|
print(f"\n预测结果: x=5 → y={model(test_x).item():.2f}")
|