PyTorch 代码流程（简单例子）

最新推荐文章于 2022-09-22 12:22:38 发布

原创

最新推荐文章于 2022-09-22 12:22:38 发布 · 1.3k 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#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