Pytorch study notes 2: training steps

Pytorch study notes 2: training steps

---

Preface

Following the previous article for the first time to implement my own data set, this article is used to introduce the pytorch training network code.

---

1. Import the data package

import cv2
import os
import torch
import numpy as np
import torchvision
import torch.nn as nn
import torch.utils.data as Data
from PIL import Image
from matplotlib.pyplot import subplot
from DLP import DLP
import torch.optim as optim
import torchvision.transforms as transforms

Two, define the data set

class OpticalSARDataset(Data.Dataset):
    def __init__(self, data_dir, part): #保存数据路径
        pass
        
    def __len__(self):
        return len(self.labels)    

    def __getitem__(self,index):
        return image,label

Three, instantiate the data class

The code is as follows (example):

# 利用之前创建好的OpticalSARDataset类去创建数据对象
train_dataset = OpticalSARDataset(data_dir, 'train')  # 训练数据集
test_dataset = OpticalSARDataset(data_dir, 'test')  # 测试数据集

Fourth, build the network

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = DLP()  # 实例化网络
model.to(device)

Five, training

epoches = 5
lr = 0.001
criterion  = nn.MSELoss().to(device) # 使用交叉熵作为损失函数
# nn.MESLoss的输入可以使矩阵和向量，默认输出的是loos的mean
optimizer = optim.Adam(model.parameters(), lr=lr)  # 定义优化器
for epoch in range(epoches):  # # 一个epoch可以认为是一次训练循环
    model.train()   #必须要写这句
    for i, data in enumerate(train_iter):
        imgs = data["image"]
        labels = data["label"]
        images = imgs.to(device)
        labels = labels.to(device)
        output = model(images)  # 经过模型对象就产生了输出
        loss = criterion(output, labels)  # 传入的参数: 输出值(预测值), 实际值(标签)
        # backprop
        optimizer.zero_grad()  # 梯度清零
        loss.backward()  # #梯度反传
        optimizer.step()  # 保留梯度

        # loss.item
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch + 1, epoches, loss.item()))
print('Finished Training')

6. Validation model

running_loss = 0
total = 0
model.eval()    #这句话也是必须的
with torch.no_grad():
    for data in test_iter:
        imgs = data["image"]
        labels = data["label"]
        images = imgs.to(device)
        labels = labels.to(device)
        output = model(images)  # 经过模型对象就产生了输出
        loss = criterion(output, labels) # 传入的参数: 输出值(预测值), 实际值(标签)
        total += labels.size(0)
        running_loss += loss.item()
        averageloss = running_loss/total

print('Loss of the network on the  test images: %f ' % (averageloss))

Seven, save the model

PATH = 'E:\pycharmProject\LaRecNet\dlp_net.pth'
torch.save(model.state_dict(), PATH)

8. Load the trained model

dlp = DLP().to(device)  # 实例化网络
dlp.load_state_dict(torch.load(PATH))
path = 'E:/pycharmProject/LaRecNet/3.jpg'
img = cv2.imread(path)
print('opencv读入图片的格式为', type(img))  # 输入图片的numpy表示
imgf = transforms.ToTensor()(img)  # 将输入图片读为向量
print(imgf.requires_grad)
imgf.requires_grad_(True)
print('转化为tensor的格式为：', type(imgf))
print('输入图片的tensor大小为：', imgf.size()) # c,h,w
imgf = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])(img).unsqueeze(0)  #把要处理的图片转化为1，c,h,w
print(imgf.shape)  # 输入图片向量的大小
out = dlp(imgf)
print(out.size())

array1 = out[0].detach().numpy()  # 将tensor数据转为numpy数据,且只要HWC
maxValue = array1.max()
array1 = array1*255/maxValue  # normalize，将图像数据扩展到[0,255]
mat = np.uint8(array1)  # float32-->uint8
print('mat_shape:', mat.shape)  # mat_shape: (3, 982, 814),tensor(3,982, 814)、numpy(982, 814, 3)存储的数据维度顺序不同
mat = mat.transpose(1, 2, 0)  # mat_shape: (982, 814，3)
cv2.imshow("feature", mat)
cv2.waitKey()

Note: There was an error when loading the saved network parameters in my place, as follows:

I did not use gpu, please refer to the description of adding a link to solve the problem

Reference link:
add link description

---