OFA モデルに学習させる - 2.5 日本語

このメソッドは、サンドイッチ ルール を使用してすべての OFA サブネットワークを一緒に最適化します。sample_random_subnet() 関数を使用してサブネットワークを生成できます。動的サブネットワークは、順方向/逆方向パスを実行できます。

各学習手順で、データの特定のミニバッチに対して、サンドイッチ ルールは ‘max’ サブネットワーク、‘min’ サブネットワーク、および 2 つのランダム サブネットワークをサンプリングします。各サブネットワークが特定のデータを使用して個別の順方向/逆方向パスを実行し、その後ですべてのサブネットワークがパラメーターをまとめてアップデートします。

# 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()