【新手入门】flux编写的神经网络无法梯度下降

``````using Flux, CUDA, Statistics, ProgressMeter
using GLMakie

## 超参数定义
n_state = 1     # 状态空间维度
n_action = 1    # 动作空间维度
n_reward = 1    # 奖励空间维度（一定是1，返回一个标量）
n_hidden = 128  # 隐藏层大小

## 创建状态价值神经网络
ValueNetwork = Chain(
Dense(n_state => n_hidden, relu; init=Flux.glorot_uniform),
BatchNorm(n_hidden), # 对每一层进行归一化，这对提升模型的泛用性和性能非常重要
Dense(n_hidden => n_hidden, relu; init=Flux.glorot_uniform),
BatchNorm(n_hidden),
Dense(n_hidden => n_reward; init=Flux.glorot_uniform),
)

## 先训练拟合一个简单线性函数作为测试
# 创建训练集
train_size = 101    # 训练集的大小
x_train = collect(range(0.0, 1.0, train_size))
x_train = reshape(x_train, (1, train_size))
y_train = 1 .- x_train
y_train

## 训练神经网络
epoches = 1000  # 定义训练次数
learning_rate = 0.005   # 定义学习率
# julia的损失函数似乎无法像pytorch那样单独提取出来
optim = Flux.setup(Flux.Adam(learning_rate), ValueNetwork)  # 为模型定义优化器
for epoch in collect(1:epoches)
y_hat = ValueNetwork(x_train)
Flux.mse(y_hat, y_train)    # 计算损失函数
end
if epoch%100 == 0
println(loss)
end
Flux.update!(optim, ValueNetwork, grads[1]) # 训练神经网络
end
``````

``````y_hat=ValueNetwork(x_train)
``````

``````y_hat=m(x_train)
``````

1 个赞