diff --git a/easybuild/easyconfigs/t/tomoDRGN/tomoDRGN-1.0.0-foss-2023a-CUDA-12.1.1.eb b/easybuild/easyconfigs/t/tomoDRGN/tomoDRGN-1.0.0-foss-2023a-CUDA-12.1.1.eb new file mode 100644 index 00000000000..b3366556457 --- /dev/null +++ b/easybuild/easyconfigs/t/tomoDRGN/tomoDRGN-1.0.0-foss-2023a-CUDA-12.1.1.eb @@ -0,0 +1,92 @@ +# Thomas Hoffmann, EMBL Heidelber, structures-it@embl.de, 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', '0.7.3'), +] + +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 = 'chem' diff --git a/easybuild/easyconfigs/t/tomoDRGN/tomoDRGN-1.0.0_use-PyTorch-2.1.patch b/easybuild/easyconfigs/t/tomoDRGN/tomoDRGN-1.0.0_use-PyTorch-2.1.patch new file mode 100644 index 00000000000..5a4ea8010f3 --- /dev/null +++ b/easybuild/easyconfigs/t/tomoDRGN/tomoDRGN-1.0.0_use-PyTorch-2.1.patch @@ -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')