Pytorch训练流程图

一 读取图片

imageFolder and DataLoader
DataLoader可自定义，但需实现两个抽象函数，分别是__len__(selft)和__getitem__(self, index)

二 定义模型

# Download and load the pretrained ResNet-18.
resnet = mm.resnet18(pretrained=True)

三 保存模型

# -------- 7. Save and load the model. --------#
# Save and load the entire model.
torch.save(resnet, 'model.mkl')
model = torch.load('model.mkl')

# Save and load only the model parameters (recommended).
torch.save(resnet.state_dict(), 'params.mkl')
resnet.load_state_dict(torch.load('params.mkl'))