pytorch中的代码封装整理

接下来大部分的blog都会是代码的整理

数据集的实现：

import os
import random
from PIL import Image
from torch.utils.data import Dataset
from torch import nn

# 主要要记住代码的主要结构
# 设定标签：
Dog_Cat_label = {"dog":0, "cat":1}
# 注意要继承Dataset父类
class DogCatDataset(Dataset):
    # 主要目的：进行transform的定义和data_info的获得
    def __init__(self,data_dir,transform=None):
        # dataset 有transform方法
        self.transform = transform
        # 在类内进行方法调用
        self.data_info = self.get_img_info(data_dir)
    # 在DataLoader中通过index读取样本
    def __getitem__(self,index):
        # index是Pytorch传入的吧，就是dataloader调用这个类的时候传入的参数
        # 一定要注意这个index，得到data_info之后注意一下
        path_img, label = self.data_info[index]
        # 直接打开图片
        img = Image.open(path_img).convert('RGB')
        # 看看有没有transform
        if self.transform is not None:
            # 这里做transform
            img = self.transform(img)
        # 统一返回img和label
        return img,label
    
    # len不要忘记写了
    def __len_(self):
        return len(self.data_info)
    
    def get_img_info(data_dir):
        data_info = list()
        # 返回根目录，文件夹和文件夹中的文件，
        # 遍历top给定的根目录下的所有文件夹，每个文件夹返回一个三元组
        # (该文件夹路径，该文件夹下的文件夹名字，该文件夹下的文件名称（不包含文件夹）)
        for root, dirs, _ in os.walk(data_dir):
            # 遍历类别
            for sub_dir in dirs:
                # 把每个subdir中的文件遍历出来，因为sub_dir就代表了图片的类别
                img_names = os.listdir(os.path.join(root,sub_dir))
                # 得到的是所有.jpg结尾的文件
                img_names = list(filter(lambda x:x.endswith('.jpg'), img_names))
                for i in range(len(img_names)):
                    # 一个一个循环遍历出图片
                    img_name = img_names[i]
                    # 得到每一个img的具体路径
                    path_img = os.path.join(root,sub_dir, img_name)
                    # 得到图片的标签,从图片里面的那个字典获取
                    label = Dog_Cat_label[sub_dir]
                    # 把整个元组返回
                    data_info.appendn((path_img,int(label)))
        return data_info

dataset总步骤

1. 首先进行img_info的获取，从get_img_info中获取每个img的标签
2. 完成__init__()中的transform和data_info的获取（当然这一部分可以在下面完成也没有事情）
3. 完成__len__()中的长度获取，比如说能够把data_info的长度
4. 完成__getitem__(self, index)函数，在里面记得传入index参数，index参数是通过DataLoader传进来的。之后对打开img，进行transform,最后返回img和标签。打开img的时候注意conver一下RGB。

模型保存与加载：

class LeNet2(nn.Module):
    def __init__(self, classes):
        super(LeNet2, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size()[0], -1)
        x = self.classifier(x)
        return x

# 先初始化模型，这个模型也比较。。。。
net = LeNet2(classes=2019)
# 记录一下保存的方法：

# 保存整个模型
torch.save(net, path_model)

# 保存模型参数，从net里面调用出state_dict()
net_state_dict = net.state_dict()
torch.save(net_state_dict, path_state_dict)

# 调用模型：这个是整个模型调用，这是torch不太赞成的一种方式
path_model = "./model.pkl"
net_load = torch.load(path_model)
print(net_load)

# 
net_new = LeNet2(classes=2019)
print("加载前: ", net_new.features[0].weight[0, ...])
# 只加载state_dict,这种其实是比较没有用的方法，因为没有办法断点续train
net_new.load_state_dict(state_dict_load)
print("加载后: ", net_new.features[0].weight[0, ...])

可以续train的save和load

# 通过一个字典进行保存，把optimizer的参数也一起保存了，这样继续训练的时候就更快了，直接无缝衔接
checkpoint = {"model_state_dict": net.state_dict(),
                      "optimizer_state_dic": optimizer.state_dict(),
                      "loss": loss,
                      "epoch": epoch}
path_checkpoint = "./checkpint_{}_epoch.pkl".format(epoch)
# 直接save
torch.save(checkpoint, path_checkpoint)

# 先加载进来字典
checkpoint = torch.load(path_checkpoint)
# 加载state_dict
net.load_state_dict(checkpoint['model_state_dict'])
# 把optimizer加载进来
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
# 把epoch加载进来
start_epoch = checkpoint['epoch']
# 无缝衔接
scheduler.last_epoch = start_epoch

百度面试题，注意一下map和reduce的使用

from functools import reduce
def func_map(x):
    return 2*x

def func_reduce(x,y):
    return 2*x+y

a = [1,3,5,7]
# 函数将一个数据集合（链表，元组等）中的所有数据进行下列操作：
# 用传给 reduce 中的函数 function（有两个参数）先对集合中的
# 第 1、2 个元素进行操作，得到的结果再与第三个数据用 function 函数
# 运算，最后得到一个结果。就相当于递归连续计算

# map第一个参数 function 以参数序列中的每一个元素调用
# function 函数，返回包含每次 function 函数返回值的新列表。
print(list(map(func_map,a)),reduce(func_reduce,a))

迁移学习，只训练功能头，不训练特征提取器

以VGG16为例

for param in my_vgg16.parameters():
    param.requires_grad = False
# 先print一下网络信息，然后保持backbone（feature）的序列不变
# 然后定义一下那些grad为true
my_vgg16.classifier[3].requires_grad = True
# 或者可以定义学习率参数组，特征提取器学习率小一些

CUDA USE

# 最重要的就是这句
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

x_cpu = torch.ones((3, 3))
print("x_cpu:\ndevice: {} is_cuda: {} id: {}".format(x_cpu.device, x_cpu.is_cuda, id(x_cpu)))
# 张量变成gpu上相当于复制一份到gpu上
# 所以必须要=，它不是原位操作
x_gpu = x_cpu.to(device)
print("x_gpu:\ndevice: {} is_cuda: {} id: {}".format(x_gpu.device, x_gpu.is_cuda, id(x_gpu)))

# 弃用，这个方法比较老了
# x_gpu = x_cpu.cuda()

net = nn.Sequential(nn.Linear(3, 3))
print("\nid:{} is_cuda: {}".format(id(net), next(net.parameters()).is_cuda))
# 直接就可以了，因为模型是原位操作的
net.to(device)
print("\nid:{} is_cuda: {}".format(id(net), next(net.parameters()).is_cuda))
# 直接把东西输入到网络里就可以了，但是呀，要注意输入的要是GPU的数据
output = net(x_gpu)

# 多GPU运行的时候要除去有序字典最开始的module.
from collections import OrderedDict
new_state_dict = OrderedDict()
# 防止多GPU运行回来报错，因为多了一个module，所以需要用这些代码使得报错小时
for k, v in state_dict_load.items():
    # 就是把module.删除，刚好7个str
    namekey = k[7:] if k.startswith('module.') else k
    new_state_dict[namekey] = v
print("new_state_dict:\n{}".format(new_state_dict))
net.load_state_dict(new_state_dict)

# 把GPU设置成python可见，这样把一些GPU留下来给别人用
os.environ.setdefault("CUDA_VISIBLE_DEVICES", gpu_list_str)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")