俗话说得好：“如果看不懂别人的代码，那一定是别人的代码写得不好！”
既然公开了代码就要对代码的可读性和正确性负责，但是看了很多源码，总是会遇到一些坑，比如：环境配置没讲清楚 -> （好不容易跑通了） 效果并没有论文中说的那么好-> （想分析代码吧）代码写得乱七八糟看不懂。。。
为了让小白能够快速上手，代码肯定需要由总到分叙述，就像讲故事一样，循序渐进才能够一目了然。下面的模板是我根据自己的理解总结的网络训练最基本的模板，大家根据自己需要再添加：
基于Pytorch的网络训练代码模板
1. 构造主函数
if __name__ == "__main__":
# 使用GPU训练
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# 模型训练的基本预设参数
batch_size = 16 # Batch Size
num_epochs = 100 # 训练迭代次数
learning_rate = 0.0001 # 学习率
root = 'data/train' # 数据集位置
##### 导入数据集
train_loader, val_loader = get_loader(root, batch_size, shuffle=True)
# 搭建模型框架
model = Model(device).to(device)
# 开始训练
train(device, model, num_epochs, learning_rate, train_loader, val_loader)
下面将具体讲解：
使用GPU训练
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
模型训练的基本预设参数
batch_size = 16 # Batch Size
num_epochs = 100 # 训练迭代次数
learning_rate = 0.0001 # 学习率
root = 'data/train' # 数据集位置
导入数据集
train_loader, val_loader = get_loader(root, batch_size, shuffle=True)
搭建模型框架
model = Model(device).to(device)
开始训练
train(device, model, num_epochs, learning_rate, train_loader, val_loader)
2.1 模型搭建
写编码器和解码器
class Model(nn.Module):
def __init__(self, device):
super(Model, self).__init__()#super的目的在于继承nn.Module并使用__init__初始化了nn.Module里的参数
self.encoder = Encoder(device) # 编码器
self.decoder = Decoder(device) # 解码器
def forward(self, x):
x = x.float()
p1 = self.encoder(x)
p2 = self.decoder(p1)
return p2
2.2 开始训练
def train(device, model, num_epochs, learning_rate, train_loader, val_loader):
print('start training ...........')
#优化器设置
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
#损失函数设置
criterion = LossCalculation(device)
train_losses, val_losses = [], []
#迭代训练
for epoch in range(epochs):
#训练集
train_epoch_loss = fit(epoch, model, optimizer, criterion, device, train_loader, phase='training')
#测试集
val_epoch_loss = fit(epoch, model, optimizer, criterion, device, val_loader, phase='validation')
print('-----------------------------------------')
#保存最优的训练参数
if epoch == 0 or val_epoch_loss <= np.min(val_losses):
torch.save(model.state_dict(), 'output/weight.pth')
#保存训练和测试的loss结果，为画图做准备
train_losses.append(train_epoch_loss)
val_losses.append(val_epoch_loss)
#绘制结果图
write_figures('output', train_losses, val_losses)
write_log('output', epoch, train_epoch_loss, val_epoch_loss)
2.2.1 优化器设置
根据自己需求设置Adam或SGD，学习率可以设置固定或自适应学习率
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
2.2.2 损失函数设置
criterion = LossCalculation(device)
class LossCalculation(nn.Module):
def __init__(self, device):
super(LossCalculation, self).__init__()
def forward(self, inputs, outputs, targets):
#inputs, outputs, targets 分别是输入图像、预测图像和label
batch_size, _, width, height = outputs.shape
total_loss = 0.0 # 一个epoch的损失函数计算
for b in range(batch_size):
#根据自己需要添加
total_loss += xxx
return total_loss
2.2.3 迭代训练
这里面最重要的就是自定义的fit函数
def fit(epoch, model, optimizer, criterion, device, data_loader, phase='training'):
if phase == 'training':
#启用一些特定的层（BN,Dropout)，设置为训练状态
model.train()
else:
#禁用一些特定的层（BN,Dropout)，设置为测试状态
model.eval()
running_loss = 0
for inputs, targets in tqdm(data_loader):
#我们使用gpu训练，那图像也必须输入到显存中
inputs = inputs.to(device)
targets = targets.to(device)
if phase == 'training':
#每一batch图像的梯度初始化为零，把loss涉及的权重的导数变成0
optimizer.zero_grad()
#输出预测图像
outputs = model(inputs)
else:
#强制之后的内容不进行计算图构建及梯度计算
with torch.no_grad():
outputs = model(inputs)
# 计算一个epoch中的一个batch的loss
loss = criterion(inputs, outputs, targets, separate_loss=False)
#累加一个epoch中的loss
running_loss += loss.item()
if phase == 'training':
#loss回传
loss.backward()
# 更新所有的参数
optimizer.step()
#计算一个epoch的loss
epoch_loss = running_loss / len(data_loader.dataset)#除以训练集或测试集图片数目
#打印一个epoch的训练结果
print('[%d][%s] loss: %.4f' % (epoch, phase, epoch_loss))
return epoch_loss