训练 OFA 模型 - 3.5 简体中文

此方法使用三明治规则来对所有 OFA 子网络进行联合最优化。sample_random_subnet() 函数用于获取子网络。动态子网络可在正向传播和反向传播中使用。

在每个训练步骤中，提供小批次数据的情况下，三明治规则会对一个“max”（最大）子网络、一个“min”（最小）子网络和两个随机子网络进行采样。每个子网络都会利用给定数据执行独立的前向/后向传递，并且所有子网络都会一起更新其参数。

# using sandwich rule and sampling subnet.
for i, (images, target) in enumerate(train_loader):

  images = images.cuda(non_blocking=True)
  target = target.cuda(non_blocking=True)

  # total subnets to be sampled
  optimizer.zero_grad()

  teacher_model.train()
  with torch.no_grad():
    soft_logits = teacher_model(images).detach()

  for arch_id in range(4):
    if arch_id == 0:
      model, _ = ofa_pruner.sample_subnet(ofa_model,'max')
    elif arch_id == 1:
      model, _ = ofa_pruner.sample_subnet(ofa_model,'min')
    else:
      model, _ = ofa_pruner.sample_subnet(ofa_model,'random') 

    output = model(images)

    loss = kd_loss(output, soft_logits) + cross_entropy_loss(output, target) 
    loss.backward()

  torch.nn.utils.clip_grad_value_(ofa_model.parameters(), 1.0)
  optimizer.step()
  lr_scheduler.step()