问具有Softmax的火炬消失梯度的线性回归

EN

class NeuralNetwork(nn.Module):
    def __init__(self,inputsize,outputsize):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(inputsize, outputsize),
            nn.Softmax(),
        )

        nn.init.uniform_(w,a=-1,b=1)

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

inputDim = 1        # takes variable 'x' 
outputDim = 1       # takes variable 'y'
learningRate = 0.001 
epochs = 100000
weight=torch.empty(3)
model = NeuralNetwork(inputDim, outputDim)
if torch.cuda.is_available():
    model.cuda()

criterion = torch.nn.MSELoss() 
optimizer = torch.optim.SGD(model.parameters(), lr=learningRate)

for epoch in range(epochs):
    # Converting inputs and labels to Variable
    if torch.cuda.is_available():
        inputs = Variable(torch.from_numpy(vS0).cuda().float())
        labels = Variable(torch.from_numpy(vC).cuda().float())
    else:
        inputs = Variable(torch.from_numpy(vS0).float())
        labels = Variable(torch.from_numpy(vC).float())

  
    optimizer.zero_grad()

    # get output from the model, given the inputs
    outputs = model(inputs)

    # get loss for the predicted output
    loss = criterion(outputs, labels)

    # get gradients w.r.t to parameters
    loss.backward()

    # update parameters
    optimizer.step()

    print('epoch {}, loss {}'.format(epoch, loss.item()))