Skip to content

fix(md): 修复消除 clash 时 Adam 最小化发散为 NaN 的问题 #75

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wszqkzqk wants to merge 1 commit into
base: lab/pku
Choose a base branch
from
+87 −12
Open

fix(md): 修复消除 clash 时 Adam 最小化发散为 NaN 的问题 #75

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions SPONGE/MD_core/min.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ struct MINIMIZATION_iteration
float beta1 = 0.9f;
float beta2 = 0.99f;
float epsilon = 1e-4f;
float learning_rate = 3e-4f;
void Gradient_Descent(int atom_numbers, VECTOR* crd, VECTOR* frc,
VECTOR* vel, const float* d_mass_inverse);
void Scale_Force_For_Dynamic_Dt(int atom_numbers, float* d_mass_inverse,
Expand Down
98 changes: 86 additions & 12 deletions SPONGE/MD_core/min.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,19 +29,75 @@ static __global__ void MD_Iteration_Gradient_Descent_With_Max_Move(
static __global__ void Get_Adam_Force(int atom_numbers, float* mass_inverse,
VECTOR* frc, VECTOR* vel, VECTOR* acc,
float beta1, float beta2, float epsilon,
float t)
float t, float learning_rate)
{
SIMPLE_DEVICE_FOR(i, atom_numbers)
{
VECTOR f = frc[i];
VECTOR f2 = {f.x * f.x, f.y * f.y, f.z * f.z};
vel[i] = beta1 * vel[i] + (1 - beta1) * mass_inverse[i] * f;
acc[i] = beta2 * acc[i] + (1 - beta2) * f2;
f = 1.0f / (1 - powf(beta1, t + 1.0f)) * vel[i];
f2 = 1.0f / (1 - powf(beta2, t + 1.0f)) * acc[i];
frc[i].x = f.x / (sqrtf(f2.x) + epsilon);
frc[i].y = f.y / (sqrtf(f2.y) + epsilon);
frc[i].z = f.z / (sqrtf(f2.z) + epsilon);
if (mass_inverse[i] == 0.0f)
{
frc[i].x = 0.0f;
frc[i].y = 0.0f;
frc[i].z = 0.0f;
}
else
{
const double first_bias = 1.0 - pow(static_cast<double>(beta1),
static_cast<double>(t) + 1.0);
const double second_bias_sqrt =
sqrt(1.0 - pow(static_cast<double>(beta2),
static_cast<double>(t) + 1.0));
VECTOR f = frc[i];
f.x = fminf(fmaxf(f.x, -1e10f), 1e10f);
f.y = fminf(fmaxf(f.y, -1e10f), 1e10f);
f.z = fminf(fmaxf(f.z, -1e10f), 1e10f);
VECTOR moment = vel[i];
VECTOR root = acc[i];
double m = static_cast<double>(beta1) * moment.x +
(1.0 - static_cast<double>(beta1)) * f.x;
double r = sqrt(static_cast<double>(beta2) * root.x * root.x +
(1.0 - static_cast<double>(beta2)) *
static_cast<double>(f.x) * f.x);
moment.x = static_cast<float>(m);
root.x = static_cast<float>(r);
f.x = static_cast<float>(learning_rate * (m / first_bias) /
(r / second_bias_sqrt + epsilon));
m = static_cast<double>(beta1) * moment.y +
(1.0 - static_cast<double>(beta1)) * f.y;
r = sqrt(static_cast<double>(beta2) * root.y * root.y +
(1.0 - static_cast<double>(beta2)) *
static_cast<double>(f.y) * f.y);
moment.y = static_cast<float>(m);
root.y = static_cast<float>(r);
f.y = static_cast<float>(learning_rate * (m / first_bias) /
(r / second_bias_sqrt + epsilon));
m = static_cast<double>(beta1) * moment.z +
(1.0 - static_cast<double>(beta1)) * f.z;
r = sqrt(static_cast<double>(beta2) * root.z * root.z +
(1.0 - static_cast<double>(beta2)) *
static_cast<double>(f.z) * f.z);
moment.z = static_cast<float>(m);
root.z = static_cast<float>(r);
f.z = static_cast<float>(learning_rate * (m / first_bias) /
(r / second_bias_sqrt + epsilon));
vel[i] = moment;
acc[i] = root;
frc[i] = f;
}
}
}

static __global__ void MD_Iteration_Adam_Move(const int atom_numbers,
VECTOR* crd, const VECTOR* frc,
const float max_move)
{
SIMPLE_DEVICE_FOR(i, atom_numbers)
{
VECTOR move = frc[i];
if (max_move > 0)
{
move = Make_Vector_Not_Exceed_Value(move, max_move);
}
crd[i] = crd[i] + move;
}
}

Expand Down Expand Up @@ -118,6 +174,15 @@ void MD_INFORMATION::MINIMIZATION_iteration::Initial(CONTROLLER* controller,
epsilon = atof(controller->Command("minimization", "epsilon"));
}
controller->printf(" minimization epsilon is %e\n", epsilon);

if (controller->Command_Exist("dt"))
{
controller->Check_Float(
"dt", "MD_INFORMATION::MINIMIZATION_iteration::Initial");
learning_rate = atof(controller->Command("dt"));
}
controller->printf(" minimization learning rate is %e A\n",
learning_rate);
}
else
{
Expand All @@ -142,7 +207,16 @@ void MD_INFORMATION::MINIMIZATION_iteration::Gradient_Descent(
int atom_numbers, VECTOR* crd, VECTOR* frc, VECTOR* vel,
const float* d_mass_inverse)
{
if (max_move <= 0)
if (dynamic_dt)
{
Launch_Device_Kernel(
MD_Iteration_Adam_Move,
(atom_numbers + CONTROLLER::device_max_thread - 1) /
CONTROLLER::device_max_thread,
CONTROLLER::device_max_thread, 0, NULL, atom_numbers, crd, frc,
max_move);
}
else if (max_move <= 0)
{
Launch_Device_Kernel(
MD_Iteration_Gradient_Descent,
Expand Down Expand Up @@ -174,6 +248,6 @@ void MD_INFORMATION::MINIMIZATION_iteration::Scale_Force_For_Dynamic_Dt(
CONTROLLER::device_max_thread,
CONTROLLER::device_max_thread, 0, NULL, atom_numbers,
d_mass_inverse, frc, vel, acc, beta1, beta2, epsilon,
md_info->sys.steps);
md_info->sys.steps, learning_rate);
}
}