训练期间如何在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>)