小萌新在用pytorch编一个关于cnn处理图片分类问题的小模型，遇到问题了，在此请教大神！

最近在用pytorch编一个关于cnn的小模型，处理的是手写体图片的分类问题，遇到问题了，在此请教大神！

模型是这样的：在cnn使用之前，为每一个图片的每一个像素分配一个径向基函数，将原本的像素值映射成该像素的函数值。然后我们需要训练的有cnn的每一个参数，以及每一个像素上径向基函数的参数，包括径向基函数的中心值以及宽度，优化方法为梯度下降法。很简单是吧？

但由于我刚学习神经网络以及pytorch的内容，遇到了一个很久也没解开的问题。就是训练的时候cnn的参数都跟着不断优化，但是径向基函数的中心值及宽度并没有实现优化。故在此请教大神，能否帮助在下实现径向基函数的中心值以及宽度的优化？（中心值和宽度分别是center和sigma两个变量）

也算比较急了，导师要求的这个任务我很久都没解决，程序比较简单了，提前感谢大神的指点！代码如下：

import torch
import torchvision
import torch.nn as nn
import torch.utils.data as Data
import matplotlib.pyplot as plt
from torch.autograd import Variable
from math import *
import numpy as np


# 超参数定义
EPOCH = 1
LR = 0.01
BATCH_SIZE = 64


# 下载MNIST数据集
train_data = torchvision.datasets.MNIST(
    root='./mnist/',
    # 是否是训练数据
    train=True,
    # 数据变换(0, 255) -> (0, 1)
    transform=torchvision.transforms.ToTensor(),
    # 是否下载MNIST数据
    download=False
)

'''
# plot one example
plt.imshow(train_data.train_data[59999].numpy(), cmap='gray')
plt.title('%i' % train_data.train_labels[4])
plt.show()
'''

test_data = torchvision.datasets.MNIST(
    root='./mnist/',
    # 是否是训练数据
    train=False,
    # 数据变换(0, 255) -> (0, 1)
    transform=torchvision.transforms.ToTensor(),
    # 是否下载MNIST数据
    download=False
)


# 数据加载
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
test_loader = Data.DataLoader(dataset=test_data, batch_size=BATCH_SIZE, shuffle=False, num_workers=1)


# 定义卷积神经网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels=1,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2
            ),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        # conv1输出为(16, 14, 14)
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # conv2输出为(32, 7, 7)
        self.output = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x, center_var, sigma_var):
        for i in range(0, 63):
            for j in range(0, 27):
                for m in range(0, 27):
                    x[i,0,j,m]=exp(-1.0*(x[i,0,j,m]-center_var[i,0,j,m])*(x[i,0,j,m]-center_var[i,0,j,m])/(2*sigma_var[i,0,j,m]*sigma_var[i,0,j,m]))
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        prediction = self.output(x)
        return prediction


cnn = CNN()

# 定义优化器
params = filter(lambda p: p.requires_grad, cnn.parameters())
optimizer = torch.optim.Adam(params, lr=LR, betas=(0.9, 0.999))

# 定义损失函数
loss_func = nn.CrossEntropyLoss()

# 训练
if __name__ == '__main__':
    for epoch in range(EPOCH):
        array = [[[[1.0] * 28] * 28] * 1] * 64
        center = np.array(array)
        sigma = np.array(array)
        center = torch.from_numpy(center)
        sigma = torch.from_numpy(sigma)
        sigma = sigma.float()
        center = center.float()
        center_var = torch.autograd.Variable(center, requires_grad=True)
        sigma_var = torch.autograd.Variable(sigma, requires_grad=True)
        for step, (x, y) in enumerate(train_loader):

            x_var = Variable(x)
            y_var = Variable(y)
            prediction = cnn(x_var, center_var, sigma_var)
            loss = loss_func(prediction, y_var)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if step % 100 == 0:
                correct = 0.0
                t=0
                for step_test, (test_x, test_y) in enumerate(test_loader):
                    test_x = Variable(test_x)
                    test_output = cnn(test_x, center_var, sigma_var)
                    pred_y = torch.max(test_output, 1)[1].data.squeeze()
                    correct += sum(pred_y == test_y)
                    t = t+1
                    if t == 156:
                        break
                accuracy = correct / 9984
                print('Epoch: ', epoch, '| train loss: %.4f' % loss.data[0], '| accuracy: ', accuracy)