Skip to content
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.

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

{bio}[foss/2023a] tomoDRGN v1.0.0, healpy v1.16.6 w/ CUDA 12.1.1 #22042

Open
wants to merge 5 commits into
base: develop
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
44 changes: 44 additions & 0 deletions easybuild/easyconfigs/h/healpy/healpy-1.16.6-foss-2023a.eb
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
easyblock = 'PythonBundle'

name = 'healpy'
version = '1.16.6'

homepage = 'http://github.com/healpy'
description = """healpy is a Python package to handle pixelated data on the sphere. It is based
on the Hierarchical Equal Area isoLatitude Pixelization (HEALPix) scheme and
bundles the HEALPix C++ library.

HEALPix was developed to efficiently process Cosmic Microwave Background data
from Cosmology experiments like BOOMERANG and WMAP but it is now used in other
branches of Astrophysics to store data from all-sky surveys. The target audience
used to be primarily the Cosmology scientific community but currently anyone
interested in handling pixelated data on the sphere is very welcome to propose
new features."""

toolchain = {'name': 'foss', 'version': '2023a'}

builddependencies = [
('cURL', '8.0.1'),
('Cython', '3.0.7'),
('poetry', '1.5.1'),
]

dependencies = [
('Python', '3.11.3'),
('astropy', '6.0.1'),
('CFITSIO', '4.3.0'),
('SciPy-bundle', '2023.07'),
('matplotlib', '3.7.2'),
]

use_pip = True
sanity_pip_check = True

exts_list = [
(name, version, {
'preinstallopts': 'sed -i "/.*define_macros/d" setup.py &&',
'checksums': ['0ab26e828fcd251a141095af6d9bf3dba43cec6f0f5cd48b65bf0af8f56329f1'],
}),
]

moduleclass = 'lib'
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
# Thomas Hoffmann, EMBL Heidelber, [email protected], 2024/12
easyblock = 'PythonBundle'

name = 'tomoDRGN'
version = '1.0.0'
versionsuffix = '-CUDA-%(cudaver)s'

homepage = 'https://github.com/bpowell122/tomodrgn/'
description = """CryoDRGN has proven a powerful deep learning method for heterogeneity analysis
in single particle cryo-EM. In particular, the method models a continuous
distribution over 3D structures by using a Variational Auto-Encoder (VAE) based
architecture to generate a reconstruction voxel-by-voxel once given a fixed
coordinate from a continuous learned latent space.

TomoDRGN extends the cryoDRGN framework to cryo-ET by learning heterogeneity
from datasets in which each particle is sampled by multiple projection images at
different stage tilt angles. For cryo-ET samples imaging particles in situ,
tomoDRGN therefore enables continuous heterogeneity analysis at a single
particle level within the native cellular environment. This new type of input
necessitates modification of the cryoDRGN architecture, enables tomography-
specific processing opportunities (e.g. dose weighting for loss weighting and
efficient voxel subset evaluation during training), and benefits from
tomography-specific interactive visualizations.
"""

toolchain = {'name': 'foss', 'version': '2023a'}

builddependencies = [
('nodejs', '18.17.1'),
('poetry', '1.5.1')
]
dependencies = [
('Python', '3.11.3'),
('CUDA', '12.1.1', '', SYSTEM),
('JupyterNotebook', '7.0.2'),
('SciPy-bundle', '2023.07'),
('PyTorch', '2.1.2', versionsuffix),
('scikit-learn', '1.3.1'),
('matplotlib', '3.7.2'),
('umap-learn', '0.5.5'),
('plotly.py', '5.16.0'),
('healpy', '1.16.6'),
('Seaborn', '0.13.2'),
('IPython', '8.14.0'),
('einops', '0.7.0'),
('adjustText', '1.1.1'),
]

use_pip = True
sanity_pip_check = True

exts_list = [
('fastcluster', '1.2.6', {
'checksums': ['aab886efa7b6bba7ac124f4498153d053e5a08b822d2254926b7206cdf5a8aa6'],
}),
('importlib_resources', '5.12.0', {
'checksums': ['4be82589bf5c1d7999aedf2a45159d10cb3ca4f19b2271f8792bc8e6da7b22f6'],
}),
('torchinfo', '1.8.0', {
'checksums': ['72e94b0e9a3e64dc583a8e5b7940b8938a1ac0f033f795457f27e6f4e7afa2e9'],
}),
(name, version, {
'patches': ['%(name)s-%(version)s_use-PyTorch-2.1.patch'],
'preinstallopts': "export SETUPTOOLS_SCM_PRETEND_VERSION=%(version)s &&",
'source_urls': ['https://github.com/bpowell122/tomodrgn/archive/refs/tags/'],
'sources': {'download_filename': 'v%(version)s.tar.gz', 'filename': '%(name)s-%(version)s.tar.gz'},
'checksums': [
{'tomoDRGN-1.0.0.tar.gz': '8b243d17af655c5088b30b4a9476072eb79e144b0a9ec1be549457c11abdd2d7'},
{'tomoDRGN-1.0.0_use-PyTorch-2.1.patch':
'5c7c1bec937ca245c999d2815f9345882950ae3da23432704efb7fb3e61f512b'},
],
}),
]

postinstallcmds = [
# required?
'mkdir %(installdir)s/share/jupyter/labextensions -p',
'export JUPYTERLAB_DIR=%(installdir)s/share/jupyter/labextensions '
'&& jupyter labextension install jupyterlab-datawidgets'
]


sanity_check_paths = {
'files': ['bin/%(namelower)s'],
'dirs': ['share/jupyter/labextensions']
}

modextravars = {'JUPYTERLAB_DIR': '%(installdir)s/share/jupyter/labextensions'} # required?

sanity_check_commands = ['%(namelower)s -h']

moduleclass = 'bio'
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
# Downgrade to PyTorch 2.1.2 (CUDA):
# relax requirement torch >= 2.3 to 2.1
# replace any occurrence torch.amp.GradScaler by conditional torch.cuda.amp / torch.cpu.amp
diff -ru tomodrgn-1.0.0/pyproject.toml tomodrgn-1.0.0_relax_requirements/pyproject.toml
--- tomodrgn-1.0.0/pyproject.toml 2024-11-15 21:15:46.000000000 +0100
+++ tomodrgn-1.0.0_relax_requirements/pyproject.toml 2024-12-16 15:36:11.187436498 +0100
@@ -31,7 +31,7 @@
"scikit-learn",
"scipy>=1.3.1",
"seaborn",
- "torch>=2.3",
+ "torch>=2.1",
"torchinfo",
"typing_extensions>=3.7.4",
"umap-learn",
diff -ru tomodrgn-1.0.0/tomodrgn/commands/train_nn.py tomodrgn-1.0.0_relax_requirements/tomodrgn/commands/train_nn.py
--- tomodrgn-1.0.0/tomodrgn/commands/train_nn.py 2024-11-15 21:15:46.000000000 +0100
+++ tomodrgn-1.0.0_relax_requirements/tomodrgn/commands/train_nn.py 2024-12-17 10:58:56.367688173 +0100
@@ -145,7 +145,7 @@


def train_batch(model: FTPositionalDecoder | DataParallelPassthrough,
- scaler: torch.amp.GradScaler,
+ scaler: torch.cuda.amp.GradScaler,
optim: torch.optim.Optimizer,
lat: Lattice,
batch_images: torch.Tensor,
@@ -402,7 +402,7 @@
# Mixed precision training with AMP
use_amp = not args.no_amp
flog(f'AMP acceleration enabled (autocast + gradscaler) : {use_amp}')
- scaler = torch.amp.GradScaler(device=device.type, enabled=use_amp)
+ scaler = torch.cuda.amp.GradScaler(enabled=use_amp) if device.type == 'cuda' else torch.cpu.amp.GradScaler(enabled=use_amp)
if use_amp:
if not args.batch_size % 8 == 0:
flog('Warning: recommended to have batch size divisible by 8 for AMP training')
diff -ru tomodrgn-1.0.0/tomodrgn/commands/train_vae.py tomodrgn-1.0.0_relax_requirements/tomodrgn/commands/train_vae.py
--- tomodrgn-1.0.0/tomodrgn/commands/train_vae.py 2024-11-15 21:15:46.000000000 +0100
+++ tomodrgn-1.0.0_relax_requirements/tomodrgn/commands/train_vae.py 2024-12-17 10:59:45.399859026 +0100
@@ -118,7 +118,7 @@

def train_batch(*,
model: TiltSeriesHetOnlyVAE | DataParallelPassthrough,
- scaler: torch.amp.GradScaler,
+ scaler: torch.cuda.amp.GradScaler,
optim: torch.optim.Optimizer,
lat: Lattice,
batch_images: torch.Tensor,
@@ -442,7 +442,7 @@

def save_checkpoint(*,
model: TiltSeriesHetOnlyVAE | DataParallelPassthrough,
- scaler: torch.amp.GradScaler,
+ scaler: torch.cuda.amp.GradScaler,
optim: torch.optim.Optimizer,
epoch: int,
z_mu_train: np.ndarray,
@@ -692,7 +692,7 @@
# Mixed precision training with AMP
use_amp = not args.no_amp
flog(f'AMP acceleration enabled (autocast + gradscaler) : {use_amp}')
- scaler = torch.amp.GradScaler(device=device.type, enabled=use_amp)
+ scaler = torch.cuda.amp.GradScaler(enabled=use_amp) if device.type == 'cuda' else torch.cpu.amp.GradScaler(enabled=use_amp)
if use_amp:
if not args.batch_size % 8 == 0:
flog('Warning: recommended to have batch size divisible by 8 for AMP training')
Loading