Skip to content

Commit

Permalink
bugfix
Browse files Browse the repository at this point in the history
  • Loading branch information
@stefanch
stefanch committed Oct 20, 2022
1 parent 87567a6 commit 29eeea6
Show file tree
Hide file tree
Showing 6 changed files with 319 additions and 173 deletions.
11 changes: 8 additions & 3 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,12 +115,17 @@ We appreciate and welcome contributions and would like to thank the following pe
Science Advances, 3(5), e1603015 (2017)
[10.1126/sciadv.1603015](http://dx.doi.org/10.1126/sciadv.1603015)

* [2] Chmiela, S., Sauceda, H. E., Müller, K.-R., & Tkatchenko, A.,
* [2] Chmiela, S., Sauceda, H. E., Müller, K.-R., Tkatchenko, A.,
*Towards Exact Molecular Dynamics Simulations with Machine-Learned Force Fields.*
Nature Communications, 9(1), 3887 (2018)
[10.1038/s41467-018-06169-2](https://doi.org/10.1038/s41467-018-06169-2)

* [3] Chmiela, S., Sauceda, H. E., Poltavsky, I., Müller, K.-R., & Tkatchenko, A.,
* [3] Chmiela, S., Sauceda, H. E., Poltavsky, I., Müller, K.-R., Tkatchenko, A.,
*sGDML: Constructing Accurate and Data Efficient Molecular Force Fields Using Machine Learning.*
Computer Physics Communications, 240, 38-45 (2019)
[10.1016/j.cpc.2019.02.007](https://doi.org/10.1016/j.cpc.2019.02.007)
[10.1016/j.cpc.2019.02.007](https://doi.org/10.1016/j.cpc.2019.02.007)

* [4] Chmiela, S., Vassilev-Galindo, V., Unke, O. T., Kabylda, A., Sauceda, H. E., Tkatchenko, A., Müller, K.-R.,
*Accurate global machine learning force fields for molecules with hundreds of atoms*
Preprint (2022)
[arXiv:2209.14865](https://arxiv.org/abs/2209.14865)
2 changes: 1 addition & 1 deletion sgdml/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

__version__ = '0.5.3.dev0'
__version__ = '0.5.3'

MAX_PRINT_WIDTH = 100
LOG_LEVELNAME_WIDTH = 7 # do not modify
Expand Down
92 changes: 82 additions & 10 deletions sgdml/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,17 @@
else:
_has_torch = True

try:
_torch_mps_is_available = torch.backends.mps.is_available()
except AttributeError:
_torch_mps_is_available = False
_torch_mps_is_available = False

try:
_torch_cuda_is_available = torch.cuda.is_available()
except AttributeError:
_torch_cuda_is_available = False

try:
import ase
except ImportError:
Expand Down Expand Up @@ -94,10 +105,14 @@ def _print_splash(max_memory, max_processes, use_torch):
max_processes_str = '{:d} CPU(s)'.format(max_processes)
hardware_str = 'using {}, {}'.format(max_memory_str, max_processes_str)

if use_torch and _has_torch and torch.cuda.is_available():
num_gpu = torch.cuda.device_count()
if num_gpu > 0:
hardware_str += ', {:d} GPU(s)'.format(num_gpu)
if use_torch and _has_torch:

if _torch_cuda_is_available:
num_gpu = torch.cuda.device_count()
if num_gpu > 0:
hardware_str += ', {:d} GPU(s)'.format(num_gpu)
elif _torch_mps_is_available:
hardware_str += ', MPS enabled'

logo_str_split = logo_str.splitlines()
print('\n'.join(logo_str_split[:-1]))
Expand Down Expand Up @@ -126,15 +141,18 @@ def _print_splash(max_memory, max_processes, use_torch):
)


_print_billboard()


def _check_update():

try:
from urllib.request import urlopen
except ImportError:
from urllib2 import urlopen

base_url = 'http://www.quantum-machine.org/gdml/'
url = '%supdate.php?v=%s' % (base_url, __version__)
base_url = 'http://api.sgdml.org/'
url = '{}update.php?v={}'.format(base_url, __version__)

can_update, must_update = '0', '0'
latest_version = ''
Expand All @@ -148,6 +166,60 @@ def _check_update():
return can_update == '1', latest_version


def _print_billboard():

try:
from urllib.request import urlopen
except ImportError:
from urllib2 import urlopen

base_url = 'http://api.sgdml.org/'
url = '{}billboard.php'.format(base_url)

resp_str = ''
try:
response = urlopen(url, timeout=1)
resp_str = response.read().decode()
response.close()
except:
pass

bbs = None
try:
import json
bbs = json.loads(resp_str)
except:
pass

if bbs is not None:

for bb in bbs:

back_color = ui.WHITE
if bb['color'] == 'YELLOW':
back_color = ui.YELLOW
elif bb['color'] == 'GREEN':
back_color = ui.GREEN
elif bb['color'] == 'RED':
back_color = ui.RED
elif bb['color'] == 'CYAN':
back_color = ui.CYAN

print(
'\n'
+ ui.color_str(
' {} '.format(bb['title']),
fore_color=ui.BLACK,
back_color=back_color,
bold=True,
)
+ '\n'
+ '-' * MAX_PRINT_WIDTH
)

print(ui.wrap_str(bb['text'], width=MAX_PRINT_WIDTH - 2))


def _print_dataset_properties(dataset, title_str='Dataset properties'):

print(ui.color_str(title_str, bold=True))
Expand Down Expand Up @@ -1692,11 +1764,11 @@ def test(

if model['use_E']:
model['e_err'] = {
'mae': np.asscalar(e_mae),
'rmse': np.asscalar(e_rmse),
'mae': e_mae.item(),
'rmse': e_rmse.item(),
}

model['f_err'] = {'mae': np.asscalar(f_mae), 'rmse': np.asscalar(f_rmse)}
model['f_err'] = {'mae':f_mae.item(), 'rmse': f_rmse.item()}
np.savez_compressed(model_path, **model)

if is_test and model['n_test'] > 0:
Expand Down Expand Up @@ -2191,7 +2263,7 @@ def _add_argument_dir_with_file_type(parser, type, or_file=False):
# Check PyTorch GPU support.
if ('use_torch' in args and args.use_torch) or 'use_torch' not in args:
if _has_torch:
if not torch.cuda.is_available():
if not (_torch_cuda_is_available or _torch_mps_is_available):
print() # TODO: print only if log level includes warning
log.warning(
'Your PyTorch installation does not see any GPU(s) on your system and will thus run all calculations on the CPU! If this is what you want, we recommend bypassing PyTorch using \'--cpu\' for improved performance.'
Expand Down
46 changes: 34 additions & 12 deletions sgdml/predict.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,17 @@
else:
_has_torch = True

try:
_torch_mps_is_available = torch.backends.mps.is_available()
except AttributeError:
_torch_mps_is_available = False
_torch_mps_is_available = False

try:
_torch_cuda_is_available = torch.cuda.is_available()
except AttributeError:
_torch_cuda_is_available = False

import numpy as np

from . import __version__
Expand Down Expand Up @@ -367,15 +378,22 @@ def __init__(
self.torch_predict = torch.nn.DataParallel(self.torch_predict)

# Send model to device
self.torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
#self.torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
if _torch_cuda_is_available:
self.torch_device = 'cuda'
elif _torch_mps_is_available:
self.torch_device = 'mps'
else:
self.torch_device = 'cpu'

while True:
try:
self.torch_predict.to(self.torch_device)
except RuntimeError as e:
if 'out of memory' in str(e):
torch.cuda.empty_cache()

print('sending to device -> fail (trying again)')
if _torch_cuda_is_available:
torch.cuda.empty_cache()

model = self.torch_predict
if isinstance(self.torch_predict, torch.nn.DataParallel):
Expand All @@ -387,9 +405,9 @@ def __init__(
model.set_n_perm_batches(
model.get_n_perm_batches() + 1
) # uncache
self.torch_predict.to(
self.torch_device
) # try sending to device again
#self.torch_predict.to( # NOTE!
# self.torch_device
#) # try sending to device again
pass
else:
self.log.critical(
Expand Down Expand Up @@ -1120,7 +1138,7 @@ def get_GPU_batch(self):
if self.use_torch:

model = self.torch_predict
if isinstance(self.torch_predict, torch.nn.DataParallel):
if isinstance(model, torch.nn.DataParallel):
model = model.module

return model._batch_size()
Expand Down Expand Up @@ -1176,11 +1194,14 @@ def predict(self, R=None, return_E=True):
print()
os._exit(1)
else:
R_torch = torch.from_numpy(R.reshape(-1, self.n_atoms, 3)).to(
R_torch = torch.from_numpy(R.reshape(-1, self.n_atoms, 3)).type(torch.float32).to(
self.torch_device
)

E_torch_F_torch = self.torch_predict.forward(R_torch, return_E=return_E)
model = self.torch_predict
if R_torch.shape[0] < torch.cuda.device_count() and isinstance(model, torch.nn.DataParallel):
model = self.torch_predict.module
E_torch_F_torch = model.forward(R_torch, return_E=return_E)

if return_E:
E_torch, F_torch = E_torch_F_torch
Expand Down Expand Up @@ -1258,9 +1279,10 @@ def predict(self, R=None, return_E=True):
)
)

E_F *= self.std
F = E_F[:, 1:]
E = E_F[:, 0] + self.c
if R is not None: # Not in train mode. TODO: better set y_std to zero
E_F *= self.std
F = E_F[:, 1:]
E = E_F[:, 0] + self.c

ret = (F,)
if return_E:
Expand Down
Loading

0 comments on commit 29eeea6

Please sign in to comment.