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}")