一、模型构建
这个写成了类，一般要继承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()
loss.backward()
optimizer.step()
|