pytorch - 如何在训练期间在softmax上设置训练权重

Question

我正在使用 PyTorch 复制 Auto-DeepLab，但我遇到了一个问题，即我无法在 softmax 上设置架构权重（包括单元格和层）。它要么导致两次后退，要么权重没有梯度升级，而只有 softmax。

class Architect () :
    def __init__(self, model, args):
        self.network_momentum = args.momentum
        self.network_weight_decay = args.weight_decay
        self.model = model
        self.optimizer = 
            torch.optim.Adam(self.model.arch_parameters(),
                lr=args.arch_lr, betas=(0.5, 0.999), 
                weight_decay=args.arch_weight_decay)

    def step (self, input_valid, target_valid) :
        self.model.soft_parameters()
        self.optimizer.zero_grad ()
        self._backward_step(input_valid, target_valid)
        self.optimizer.step()

    def _backward_step (self, input_valid, target_valid) :
        _, loss = self.model._loss (input_valid, target_valid)
        loss.backward ()

在这段代码下，它导致了 softmax 的工作，但权重没有使用梯度进行优化。

[[0.3333, 0.3333, 0.3333],
         [0.3333, 0.3333, 0.3333],
         [0.3333, 0.3333, 0.3333],
         [0.3333, 0.3333, 0.3333]]*12], device='cuda:0',
       grad_fn=<SoftmaxBackward>)