PyTorch 代码流程（简单例子）

一、模型构建

这个写成了类，一般要继承torch.nn.Module来定义网络结构，然后再通过forward()定义前向过程。

下面以一个很简单的两层全连接网络为例：

# net
class net(nn.Module):
  def __init__(self):
    super(net, self).__init__() 
    self.fc1 = nn.Linear(50, 50)
    self.fc2 = nn.Linear(50, 10)
  def forward(self, x):
    fc1 = self.fc1(x)
    fc2 = self.fc2(fc1)
    return fc2

# model
net = net()

其中，super这个关键字主要是用于调用父类的方法，它可以防止对父类的多次调用，相当于产生了一个super类的对象。

二、数据处理

数据处理一般是用官方给的Dataset抽象类，根据数据的特点处理。也可不用官方的类，自行处理数据。还有可能是用现成的数据集。

这里是一个txt文件保存了图片路径与单个标签的例子：

from PIL import Image
from torch.utils.data import Dataset
class trainDataset(Dataset):
    def __init__(self, txt_path, transform=None, target_transform=None):
        fh = open(txt_path, 'r')
        imgs = []
        for line in fh:
            line = line.rstrip()
            words = line.split()
            imgs.append((words[0].int(words[1]))) # 图片路径+label
        self.imgs = imgs
        self.transform = transform        

    def __getitem__(self, index):
        fn, label = self.imgs[index]
        img = Image.open(fn).convert('RGB')
        if self.transform is not None:
            img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.imgs)

然后加载训练数据集与验证集，同时可以实现取batch，shuffle或者多线程读取数据，这里是最简单的版本：

# data
train_data = torch.utils.data.DataLoader(trainDataset)
val_data = torch.utils.data.DataLoader(valDataset)

当然还可以加上一些数据增强的方法。

三、定义Loss和optimizer

# loss
loss = torch.nn.CrossEntropyloss()
loss = torch.nn.MSELoss()

# optimizer
optimizer = torch.optim.Adam(net.parameters())

四、构建训练过程

以一个epoch的训练为例：

def train(epoch):
    for i, data in enumerate(dataLoader, 0):
        x, y = data
        y_pred = model(x)

        # loss compute
        loss = criterion(y_pred, y)
        optimizer.zero_grad()  # pytorch梯度会自动累加，反向传播之前需要将梯度清零
        loss.backward()
        optimizer.step()