A Simple Example of PyTorch for Initiates

Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., et al. (2019). PyTorch: An imperative style, high-performance deep learning library. ArXiv e-prints, pp. arXiv:1912.01703.

Thanks all the authors of PyTorch for their great contribution.

Here, I give a simple example of CNN implemented with PyTorch for the initiates. I hope that it can be helpful.

First, we import the library.

import torch
import torch.nn as nn
import numpy as np

Net, we build a simple dataset. Here, x = { x 0 , x 1 , x 2 , x 3 } \bm{x} = \{x_0, x_1, x_2, x_3\} x={x0​,x1​,x2​,x3​} and $x_i$ is 2-dimension. Only when $x_i^{(0)}\oplus x_i^{(1)}=1$, the label $y_i$ of $x_i$ is 1.

x = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
x = torch.Tensor(x).float()
y = np.array([1, 0, 0, 1])
y = torch.Tensor(y).long()

Ok, now let’s build a simple CNN model. Note that the activation function of the last layer is Linear not Relu. This is because if too many activations get below zero then most of the units (neurons) in network with ReLu will simply output zero, in other words, die and thereby prohibiting learning.

class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()

        self.linear1 = nn.Linear(2, 50)
        self.linear2 = nn.Linear(50, 30)
        self.linear3 = nn.Linear(30, 2)

        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.dropout = nn.Dropout(p=0.5)

    def forward(self, x):
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        x = self.dropout(x)
        x = self.linear3(x)

        return x

Then, we begin to train the network.

net = MyNet()
# train() enables some modules like dropout, and eval() does the opposit
net.train()

# set the optimizer where lr is the learning-rate
optimizer = torch.optim.SGD(net.parameters(), lr=0.05)
loss_func = nn.CrossEntropyLoss()

for epoch in range(50000):
    if epoch % 5000 == 0:
    	# call eval() and evaluate the model on the validation set
    	# when calculate the loss value or evaluate the model on the validation set,
    	# it's suggested to use "with torch.no_grad()" to save the memory. Here I didn't use it.
        net.eval()
        out = net(x)
        loss = loss_func(out, y)
        print(loss.detach().numpy())
        # call train() and train the model on the training set
        net.train()

    out = net(x)
    loss = loss_func(out, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 5000 == 0:
        net.eval()
        out = net(x)
        loss = loss_func(out, y)
        print(loss.detach().numpy())
        print('----')
        net.train()

    if epoch % 1000 == 0:
    	# adjust the learning-rate
    	# weight decay every 1000 epochs
        lr = optimizer.param_groups[0]['lr']
        lr *= 0.9
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

Test the trained model.

net.eval()
print(net(x).data)

All right, let’s save the model.

path = './save/state_dict_model.pth'
torch.save(net.state_dict(), path)

Load the model.

model = MyNet()
model.load_state_dict(torch.load(path))
# disable the dropout
model.eval()

print(model(x).data)
# print the label
print(np.argmax((model(x).detach().numpy()), axis=1))

Here, I give another example about $L_p$ regularization. In PyTorch, we can use

optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=0.0005)

where

weight_decay (float, optional): weight decay (L2 penalty) (default: 0)

to enable the $L_2$ regularization. However, if we want to use $L_1$ regularization, we must recode it ourselves.

import torch

class Regularization(torch.nn.Module):
    def __init__(self, model, weight_decay, p=2):
        '''
        :param model 模型
        :param weight_decay:正则化参数
        :param p: 范数计算中的幂指数值，默认求2范数,
        '''
        super(Regularization, self).__init__()
        if weight_decay <= 0:
            raise Exception('the value of weight_decay must be larger than 0.')
        self.model = model
        self.weight_decay = weight_decay
        self.p = p
        self.weight_list = self.get_weight(model)

    def forward(self):
        self.weight_list = self.get_weight(self.model)  # 获得最新的权重
        reg_loss = self.regularization_loss(self.weight_list, self.weight_decay, p=self.p)
        return reg_loss

    def get_weight(self, model):
        '''
        获得模型的权重列表
        :param model:
        :return:
        '''
        weight_list = []
        for name, param in model.named_parameters():
        	# exclude the bias
            if 'weight' in name:
                weight = (name, param)
                weight_list.append(weight)
        return weight_list

    def regularization_loss(self, weight_list, weight_decay, p=2):
        '''
        计算张量范数
        :param weight_list:
        :param p: 范数计算中的幂指数值，默认求2范数
        :param weight_decay:
        :return:
        '''
        reg_loss = 0
        for name, w in weight_list:
            l2_reg = torch.norm(w, p=p)
            reg_loss = reg_loss + l2_reg

        reg_loss = weight_decay * reg_loss
        return reg_loss

Use the $L_p$ regularization.

import torch
import torch.nn as nn
import numpy as np

from regularization import Regularization

# 构建输入集
x = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
x = torch.Tensor(x).float()
y = np.array([1, 0, 0, 1])
y = torch.Tensor(y).long()


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()

        self.linear1 = nn.Linear(2, 50)
        self.linear2 = nn.Linear(50, 30)
        self.linear3 = nn.Linear(30, 2)

        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.dropout = nn.Dropout(p=0.5)

    def forward(self, x):
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        x = self.dropout(x)
        x = self.linear3(x)

        return x


net = MyNet()
net.train()

# 设置优化器
optimizer = torch.optim.SGD(net.parameters(), lr=0.05)
loss_func = nn.CrossEntropyLoss()

loss_l2_func = Regularization(net, weight_decay=0.0005, p=2)

for epoch in range(50000):
    out = net(x)
    loss_cls = loss_func(out, y)
    # loss_l2_func() calls the Regularization.forward() function to update the values of the weights
    loss_total = loss_cls + loss_l2_func()
    optimizer.zero_grad()
    loss_total.backward()
    optimizer.step()

    if epoch % 1000 == 0:
        lr = optimizer.param_groups[0]['lr']
        lr *= 0.9
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

net.eval()
print(net(x).data)

path = './save/state_dict_model.pth'

torch.save(net.state_dict(), path)

model = MyNet()
# 载入保存的模型参数
model.load_state_dict(torch.load(path))
# 不启用Dropout
model.eval()

print(model(x).data)
print(np.argmax((model(x).detach().numpy()), axis=1))

At last, again, I appreciate the authors of PyTorch very much for their invaluable contribution. Besides, I am extremely grateful to the open source code contributors.

Reference

1. （深度学习）Pytorch之dropout训练
2. PyTorch中的学习率调整函数
3. python日记：用pytorch搭建一个简单的神经网络
4. 【PyTorch】训练一个最简单的CNN
5. pytorch实现L2和L1正则化regularization的方法
6. PyTorch实现一个简单的神经网络 （回归模型）