pytorch保存提取

代码部分

import torch

import matplotlib.pyplot as plt

 

# torch.manual_seed(1) # reproducible

 

# fake data

x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1) # x data (tensor), shape=(100, 1)

y = x.pow(2) + 0.2*torch.rand(x.size()) # noisy y data (tensor), shape=(100, 1)

 

# The code below is deprecated in Pytorch 0.4. Now, autograd directly supports tensors

# x, y = Variable(x, requires_grad=False), Variable(y, requires_grad=False)

 

 

def save(q,w,e):

    # save net1

    net1 = torch.nn.Sequential(

        torch.nn.Linear(q,w),

        torch.nn.ReLU(),

        torch.nn.Linear(w,e)

    )

    optimizer = torch.optim.SGD(net1.parameters(), lr=0.5)

    loss_func = torch.nn.MSELoss()

 

    for t in range(100):

        prediction = net1(x)

        loss = loss_func(prediction, y)

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

 

    # plot result

    plt.figure(1, figsize=(10, 3))

    plt.subplot(131)

    plt.title('Net1')

    plt.scatter(x.data.numpy(), y.data.numpy())

    plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)

 

    # 2 ways to save the net

    torch.save(net1, 'net.pkl') # save entire net

    torch.save(net1.state_dict(), 'net_params.pkl') # save only the parameters

 

 

def restore_net():

    # restore entire net1 to net2

    net2 = torch.load('net.pkl')

    prediction = net2(x)

 

    # plot result

    plt.subplot(132)

    plt.title('Net2')

    plt.scatter(x.data.numpy(), y.data.numpy())

    plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)

 

 

def restore_params(q,w,e):

    # restore only the parameters in net1 to net3

    net3 = torch.nn.Sequential(

        torch.nn.Linear(q,w),

        torch.nn.ReLU(),

        torch.nn.Linear(w,e)

    )

 

    # copy net1's parameters into net3

    net3.load_state_dict(torch.load('net_params.pkl'))

    prediction = net3(x)

 

    # plot result

    plt.subplot(133)

    plt.title('Net3')

    plt.scatter(x.data.numpy(), y.data.numpy())

    plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)

    plt.show()

 

# save net1

save(1,10,1)

 

# restore entire net (may slow)

restore_net()

 

# restore only the net parameters

restore_params(1,10,1)

运行截图

 改地址