重构代码以避免训练循环中的 for 循环？

Question

我正在定义一个火车函数，我将其传递给 data_loader 作为字典。

• data_loader['train']：由训练数据组成
• data_loader['val'] 由验证数据组成。

我创建了一个循环，它迭代我所处的阶段（train 或 val），并将模型相应地设置为 model.train() 或 model.eval() 。但是我觉得这里有太多嵌套的 for 循环，这使得计算成本很高。有人能推荐一种更好的方法来构建我的火车功能吗？我应该创建一个单独的函数来进行验证吗？

以下是我到目前为止所拥有的：

#Make train function (simple at first)
def train_network(model, optimizer, data_loader, no_epochs):

  total_epochs = notebook.tqdm(range(no_epochs))

  for epoch in total_epochs:
    
    for phase in ['train', 'val']:
      if phase == 'train':
        model.train()
      else:
        model.eval()

      for i, (images, g_truth) in enumerate(data_loader[phase]):
        images = images.to(device)
        g_truth = g_truth.to(device)

Answer 1

在编写训练脚本时，最外层和最内层的 for 循环很常见。

我看到的最常见的模式是：

total_epochs = notebook.tqdm(range(no_epochs))

for epoch in total_epochs:
    # Training
    for i, (images, g_truth) in enumerate(train_data_loader):
        model.train()
        images = images.to(device)
        g_truth = g_truth.to(device)
        ...

    # Validating
    for i, (images, g_truth) in enumerate(val_data_loader):
        model.eval()
        images = images.to(device)
        g_truth = g_truth.to(device)
        ...

如果您需要使用以前的变量 data_loader，则可以将 train_data_loader 替换为 data_loader["train"] 和 val_data_loader 和 data_loader["val"]

这种布局很常见，因为我们通常希望在验证时而不是训练时做一些不同的事情。这也可以更好地构建代码，并避免在最内层循环的不同部分可能需要的大量 if stage == "train" 。但这确实意味着您可能需要重复一些代码。这种权衡是普遍接受的，如果我们有 3 个或更多阶段（例如多个验证阶段或评估阶段），则可能会考虑您的原始代码。