I am trying to make a regression with NN using Flux. I have separated the dataset in the train/test parts and implemented a simple neural network.
x_train, x_test, y_train, y_test = train_test_split(convert(Array, x), convert(Array, y), test_size=.3)
nndata = Flux.Data.DataLoader((x_train’, y_train), batchsize=30, shuffle=false)
model = Chain(
Dense(23, 40, relu),
Dense(40, 25, relu),
Dense(25, 1, identity),
)
ps = Flux.params(model)
loss(x_train, y_train) = Flux.mse(model(x_train), y_train)
opt = ADAM()
using Flux: train!
for i in Array((1:50)’)
train!(loss, ps, nndata, opt)
end
y_pred = model(x_train’)
When I train the model, after some epochs, all the predicted value points to a single value, which is close to the mean of y_train.
Does anyone have a guess as on what is going on?
Thank you