常用pytorch代码

准备工作

导包

提前安装tensorboardX （可视化），torchsnooper（调试代码）

pip install tensorboardX
pip install torchsnooper

import torch
from tensorboardX import SummaryWriter
import torch.nn as nn
import matplotlib.pyplot as plt
import torch.nn.functional as F
import numpy as np
import torch.optim as optim
import torchsnooper
import random
from copy import deepcopy
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import math
from collections import OrderedDict
%matplotlib inline

设置随机数

为了可复现，每次运行代码时设置一次随机数。

def set_seed(seed):
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)

定义网络模型

#方式一 定义网络模型类
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(1, 20)
        self.fc2 = nn.Linear(20, 20)
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return x
 net = Net()
 
 #方式二 用nn.Sequential定义
net = nn.Sequential(nn.Linear(1, 20),nn.ReLU(),nn.Linear(20, 20),nn.ReLU())

#方式三 用add_module创建
net = nn.Sequential()
for i in range(2):
	net.add_module('fc'+str(i),nn.Linear(20, 20))
	net.add_module('relu'+str(i),nn.ReLU())

常用网络模块：

nn.Conv2d(in_channels=?, out_channels=?, kernel_size=(?, ?), stride=1, padding=0, dilation=1)
nn.MaxPool2d(kernel_size=(?,?), stride=None, padding= 0, dilation= 1, return_indices = False)
nn.BatchNorm2d(num_features=?, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
nn.RNN(input_size=?, hidden_size=?, num_layers=?,batch_first=False,dropout=0,bidirectional=False)
nn.LSTM(input_size=?, hidden_size=?, num_layers=?,batch_first=False,dropout=0,bidirectional=False)
nn.Dropout(p=0.5, inplace=False)
nn.Embedding(num_embeddings=?, embedding_dim=?, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False, _weight=None)

转化数据类型（非必须）

遇到数据类型不匹配的问题时，都可通过以下代码转化。

x.float()
x.double()
x.int()
x.long()
x.byte()

定义损失函数

loss_fn = nn.MSELoss() #均方误差
loss_fn = nn.CrossEntropyLoss() #多分类，损失函数中已包括softmax，所以网络最后一层不需要经过softmax激活函数
loss_fn = nn.BCEWithLogitsLoss()#二分类，损失函数已包括sigmoid和之后取log
loss_fn = nn.L1Loss()#绝对值损失函数

以上损失函数可以设置样本权重，是否取平均之类的，请查阅官方文档。

定义优化器

optimizer = optim.SGD(net.parameters(), lr=?)#输入lr的值 固定学习率（步长）
optimizer = optim.SGD(net.parameters(), lr=?,, momentum=?)#对学习率（步长）指数加权移动平均(momentun)
optimizer = optim.Adagrad(net.parameters(), lr=?,weight_decay=0, initial_accumulator_value=0, eps=1e-10)#步长为学习率除以梯度的平方累计和的均方根
optimizer = optim.RMSprop(net.parameters(), lr=?, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)#步长为学习率除以梯度的平方累计和的指数加权移动平均的均方根
optimizer = optim.Adadelta(net.parameters(), rho=0.9, eps=1e-06, weight_decay=0)#步长为之前步长的指数加权移动平均的均方根，除以梯度的平方累计和的指数加权移动平均的均方根
optimizer = optim.Adam(net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)#步长为学习率除以梯度的平方累计和的指数加权移动平均的均方根的无偏估计，同时乘以的梯度为指数加权移动平均的无偏估计

weight_decay 参数就是设置L2正则化系数。

动态调整学习率（非必须）

scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma, last_epoch=-1, verbose=False)#每过一步就乘以gamma
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=?, gamma=0.1, last_epoch=-1, verbose=False)#每step_size步后lr就乘以gamma
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[?], gamma=0.1, last_epoch=-1, verbose=False)#每到达milestones的步数后就乘以gamma
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, threshold=0.0001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-08, verbose=False)#min模式：metric不再下降减小学习率。max模式：不再上升减小学习率。

用法：

#如果是ReduceLROnPlateau则需要传递metric参数
scheduler.step(loss)
#其他的scheduler直接调用step函数，不用传参。
scheduler.step()

定义dataset（非必须）

class MyDataset(Dataset):
    def __init__(self,dataset,label):
        self.dataset=dataset
        self.label=label
    def __getitem__(self,index):
        element = self.dataset[index]
        label =self.label[0,index]
        return element,label
    def __len__(self):
        return self.dataset.shape[0]
train_dataloader=DataLoader(MyDataset(train_X,train_y),batch_size=4,shuffle=True)
testloader=DataLoader(MyDataset(test_X,test_y))

训练测试网络

训练网络

numOfepoches = 30
for epoch in range(numOfepoches):
    for i,data in enumerate(train_dataloader,0):
        X,y = data
        #X,y=X.float(),y.float().view(-1,1)#非必须，视自己代码情况而定
        output=net(X)        
        loss = loss_fn(output,y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        #可自己加入打印loss或者输出到tensorboard的代码
    #scheduler.step()#非必须，视自己代码情况而定

测试网络

以下代码是二分类（可以拓展到多分类）代码，网络输出没有sigmoid激活函数。

correct = 0
total = 0
zero=torch.tensor(0)
one=torch.tensor(1)
net.eval()#不要忘记写这句代码
with torch.no_grad():
    for i,data in enumerate(testloader):
        X, y= data
        X=X.float()
        output=net(X)
        predicted = torch.where(output>=0,one,zero)#大于等于零分类为1，否则为0
        #如果是多分类，上面的代码改成下面代码:
        #predicted = torch.argmax(output, dim=?, keepdim=False)
        total += y.size(0)
        correct += (predicted == y).sum().item()

print('Accuracy of the network on the test images: %d %%' % (
    100 * correct / total))

保存模型

#方式一：只保存参数
#注意使用deepcopy，否则best_params会变化
best_params =  deepcopy(net.state_dict())
torch.save(best_params, 'best_params.pkl')
#方式二：保存整个模型
torch.save(net, 'model.pkl')

加载模型

#方式一：加载参数（网络结构必须一致，包括每层大小名字，否则报错）
best_params =  torch.load('best_params.pkl')
net.load_state_dict(best_params)
#方式二：直接加载整个模型
net= torch.load('model.pkl')

获取某些变量的中间结果

#Tensor变量注册hook函数，每次反向传播后就执行一次
def print_grad(grad):
    print('grad is \n',grad)
x.register_hook(print_grad)
#Module注册函数forhook和backhook,为了传递进来name,所以外面包多了一层forwardhook和backwardhook，采用tensorboard记录输入和输出。
#之所以下面代码写那么长，是因为input和output有可能是tuple，也有可能是tensor，还有可能是None，所以做了很多判断。
writer = SummaryWriter('runs/test')
def forwardhook(name):
    def forhook(module, input, output):
        if(torch.is_tensor(input)):
            #下面这行换成你想处理的代码
            writer.add_histogram(name+"_input_1", input.clone().cpu().data.numpy(), 1)
        else:  
            for i,v in enumerate(input):
                if v is not None:
                    #下面这行换成你想处理的代码
                    writer.add_histogram(name+"_input_"+str(i), v.clone().cpu().data.numpy(), 1)
        if(torch.is_tensor(output)):
            #下面这行换成你想处理的代码
            writer.add_histogram(name+"_output_1", output.clone().cpu().data.numpy(), 1)
        else:  
            for i,v in enumerate(output):
                if v is not None:
                    #下面这行换成你想处理的代码
                    writer.add_histogram(name+"_output_"+str(i), v.clone().cpu().data.numpy(), 1)
    return forhook
#判断生成器是否为空
def IsGenEmpty(gen):
    try:
        gen.__next__()
        return False
    except:
        return True
#以上代码都是准备工作，可根据自己需要改写
forward_handles = []
for name,module in net.named_modules():
    if(IsGenEmpty(module.children())):
        #关键代码register_forward_hook
        handle = module.register_forward_hook(forwardhook(name))
        #注意：不要使用register_backward_hook，因为有bug!
        #之所以保存返回的handle，是因为之后需要注销的时候用到
        forward_handles.append(handle)
#下面可以运行网络，正向传播和反向传播，然后就会自动调用上面注册的hook函数
#do something

#最后如果需要注销hook函数需要下面代码：
for handle in forward_handles:
    handle.remove()#关键代码

其他一些代码

调试

#方式一
@torchsnooper.snoop()#将装饰器放在要调试的函数前即可
def myfunc(X,y):
	pass

#方式二
with torchsnooper.snoop():#用with语句来激活 TorchSnooper，包裹需要调试的代码
	pass

tensorboardX 代码

创建writer

#日志文件保存在 'runs/exp-1'
writer = SummaryWriter('runs/exp-1')
#日志文件默认保存在'runs/日期'，如'runs/Aug20-17-20-33'
writer2 = SummaryWriter()
#在保存文件名后加comment值，如 'runs/Aug20-17-20-33-3xlearning rate'
writer3 = SummaryWriter(comment='3x learning rate')

写入数据

for name, param in net.named_parameters():
        writer.add_histogram(name, param.clone().cpu().data.numpy(), epoch)#记录权重
        if(param.grad is not None):
            writer.add_histogram(name+"_grad", param.grad.numpy(), epoch)#记录权重的梯度
writer.add_scalar('loss', loss.item(), global_step=epoch)#记录loss

刷新数据

如果不刷新数据，可能tensorboard不会即时显示。

write.flush()

关闭写入文件

write.close()

读写图片

设置数据集需要读取图片，这里写一下常用代码。

import cv2
img_cv = cv2.imread(path)#opencv读取数据是BGR通道，高×宽×通道数
img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)#转化为RGB格式
tensor_cv = torch.from_numpy(np.transpose(img_cv, (2, 0, 1)))#pytorch处理图象是最后两维为高×宽，即样本数x通道数x高x宽
plt.imshow(img_cv)#显示图片 一般在jupyter notebook中显示