from operator import truediv
import numpy as np
import torch
from torch.utils import data
from torch import nn
import d2l
import matplotlib.pyplot as plt
n_train,n_test,num_inputs,batch_size=20,100,200,5#训练样本少，容易过拟合
true_w,true_b=torch.ones((num_inputs,1))*0.01,0.5
def synthetic_data(m,num_inputs):
    x=torch.normal(0,0.1,(m,num_inputs))
    noise = torch.normal(0, 0.5, (m, 1))
    y=torch.matmul(x,true_w)+true_b+noise
    return x,y
features_train,labels_train=synthetic_data(n_train,num_inputs)
features_test,labels_test=synthetic_data(n_test,num_inputs)
def load_array(array,batch_size,is_train=True):
    dataset=data.TensorDataset(*array)
    return data.DataLoader(dataset,batch_size,shuffle=is_train)
train_iter=load_array((features_train,labels_train),batch_size,is_train=True)
test_iter=load_array((features_test,labels_test),batch_size,is_train=False)

def init_param():
    w=torch.normal(0,0.1,(num_inputs,1),requires_grad=True)
    b=torch.zeros(1,requires_grad=True)
    return [w,b]
def l2_penalty(w):
    return torch.sum(w.pow(2))/2
def squared_loss(yhat,y):
    return (yhat-y.reshape(yhat.shape))**2 /2
def evaluate(data_iter,net,loss):
    metric,n=0.0,0
    for x,y in data_iter:
        yhat=net(x)
        l=loss(yhat,y)
        metric+=l.sum().item()
        n+=y.numel()
    return metric/n
def train(lambd):
    w,b=init_param()
    net,loss=lambda X:X@w+b,squared_loss
    num_epoch,lr=100,0.003
    animator=d2l.Animator(xlabel='epochs',ylabel='loss',legend=['train','test'],xlim=[5,num_epoch],yscale='log',nrows=1,ncols=1,figsize=(7,5))
    for epoch in range(num_epoch):
        for x,y in train_iter:
            l=loss(net(x),y)+lambd*l2_penalty(w)
            l.sum().backward()
            d2l.sgd([w,b],lr,batch_size)
        if (epoch+1)%5==0:
            train_loss,test_loss=evaluate(train_iter,net,loss),evaluate(test_iter,net,loss)
            animator.add(epoch+1,(train_loss,test_loss))

    print('w的L2范数是：', torch.norm(w).item())
train(lambd=0.5)
plt.ioff()
plt.show(block=True)