diff --git a/bitsandbytes_windows/cextension.py b/bitsandbytes_windows/cextension.py new file mode 100644 index 00000000..d38684a2 --- /dev/null +++ b/bitsandbytes_windows/cextension.py @@ -0,0 +1,54 @@ +import ctypes as ct +from pathlib import Path +from warnings import warn + +from .cuda_setup.main import evaluate_cuda_setup + + +class CUDALibrary_Singleton(object): + _instance = None + + def __init__(self): + raise RuntimeError("Call get_instance() instead") + + def initialize(self): + binary_name = evaluate_cuda_setup() + package_dir = Path(__file__).parent + binary_path = package_dir / binary_name + + if not binary_path.exists(): + print(f"CUDA SETUP: TODO: compile library for specific version: {binary_name}") + legacy_binary_name = "libbitsandbytes.so" + print(f"CUDA SETUP: Defaulting to {legacy_binary_name}...") + binary_path = package_dir / legacy_binary_name + if not binary_path.exists(): + print('CUDA SETUP: CUDA detection failed. Either CUDA driver not installed, CUDA not installed, or you have multiple conflicting CUDA libraries!') + print('CUDA SETUP: If you compiled from source, try again with `make CUDA_VERSION=DETECTED_CUDA_VERSION` for example, `make CUDA_VERSION=113`.') + raise Exception('CUDA SETUP: Setup Failed!') + # self.lib = ct.cdll.LoadLibrary(binary_path) + self.lib = ct.cdll.LoadLibrary(str(binary_path)) # $$$ + else: + print(f"CUDA SETUP: Loading binary {binary_path}...") + # self.lib = ct.cdll.LoadLibrary(binary_path) + self.lib = ct.cdll.LoadLibrary(str(binary_path)) # $$$ + + @classmethod + def get_instance(cls): + if cls._instance is None: + cls._instance = cls.__new__(cls) + cls._instance.initialize() + return cls._instance + + +lib = CUDALibrary_Singleton.get_instance().lib +try: + lib.cadam32bit_g32 + lib.get_context.restype = ct.c_void_p + lib.get_cusparse.restype = ct.c_void_p + COMPILED_WITH_CUDA = True +except AttributeError: + warn( + "The installed version of bitsandbytes was compiled without GPU support. " + "8-bit optimizers and GPU quantization are unavailable." + ) + COMPILED_WITH_CUDA = False diff --git a/bitsandbytes_windows/cuda_setup/main.py b/bitsandbytes_windows/cuda_setup/main.py new file mode 100644 index 00000000..7e5f9c98 --- /dev/null +++ b/bitsandbytes_windows/cuda_setup/main.py @@ -0,0 +1,166 @@ +""" +extract factors the build is dependent on: +[X] compute capability + [ ] TODO: Q - What if we have multiple GPUs of different makes? +- CUDA version +- Software: + - CPU-only: only CPU quantization functions (no optimizer, no matrix multiple) + - CuBLAS-LT: full-build 8-bit optimizer + - no CuBLAS-LT: no 8-bit matrix multiplication (`nomatmul`) + +evaluation: + - if paths faulty, return meaningful error + - else: + - determine CUDA version + - determine capabilities + - based on that set the default path +""" + +import ctypes + +from .paths import determine_cuda_runtime_lib_path + + +def check_cuda_result(cuda, result_val): + # 3. Check for CUDA errors + if result_val != 0: + error_str = ctypes.c_char_p() + cuda.cuGetErrorString(result_val, ctypes.byref(error_str)) + print(f"CUDA exception! Error code: {error_str.value.decode()}") + +def get_cuda_version(cuda, cudart_path): + # https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART____VERSION.html#group__CUDART____VERSION + try: + cudart = ctypes.CDLL(cudart_path) + except OSError: + # TODO: shouldn't we error or at least warn here? + print(f'ERROR: libcudart.so could not be read from path: {cudart_path}!') + return None + + version = ctypes.c_int() + check_cuda_result(cuda, cudart.cudaRuntimeGetVersion(ctypes.byref(version))) + version = int(version.value) + major = version//1000 + minor = (version-(major*1000))//10 + + if major < 11: + print('CUDA SETUP: CUDA version lower than 11 are currently not supported for LLM.int8(). You will be only to use 8-bit optimizers and quantization routines!!') + + return f'{major}{minor}' + + +def get_cuda_lib_handle(): + # 1. find libcuda.so library (GPU driver) (/usr/lib) + try: + cuda = ctypes.CDLL("libcuda.so") + except OSError: + # TODO: shouldn't we error or at least warn here? + print('CUDA SETUP: WARNING! libcuda.so not found! Do you have a CUDA driver installed? If you are on a cluster, make sure you are on a CUDA machine!') + return None + check_cuda_result(cuda, cuda.cuInit(0)) + + return cuda + + +def get_compute_capabilities(cuda): + """ + 1. find libcuda.so library (GPU driver) (/usr/lib) + init_device -> init variables -> call function by reference + 2. call extern C function to determine CC + (https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html) + 3. Check for CUDA errors + https://stackoverflow.com/questions/14038589/what-is-the-canonical-way-to-check-for-errors-using-the-cuda-runtime-api + # bits taken from https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549 + """ + + + nGpus = ctypes.c_int() + cc_major = ctypes.c_int() + cc_minor = ctypes.c_int() + + device = ctypes.c_int() + + check_cuda_result(cuda, cuda.cuDeviceGetCount(ctypes.byref(nGpus))) + ccs = [] + for i in range(nGpus.value): + check_cuda_result(cuda, cuda.cuDeviceGet(ctypes.byref(device), i)) + ref_major = ctypes.byref(cc_major) + ref_minor = ctypes.byref(cc_minor) + # 2. call extern C function to determine CC + check_cuda_result( + cuda, cuda.cuDeviceComputeCapability(ref_major, ref_minor, device) + ) + ccs.append(f"{cc_major.value}.{cc_minor.value}") + + return ccs + + +# def get_compute_capability()-> Union[List[str, ...], None]: # FIXME: error +def get_compute_capability(cuda): + """ + Extracts the highest compute capbility from all available GPUs, as compute + capabilities are downwards compatible. If no GPUs are detected, it returns + None. + """ + ccs = get_compute_capabilities(cuda) + if ccs is not None: + # TODO: handle different compute capabilities; for now, take the max + return ccs[-1] + return None + + +def evaluate_cuda_setup(): + print('') + print('='*35 + 'BUG REPORT' + '='*35) + print('Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues') + print('For effortless bug reporting copy-paste your error into this form: https://docs.google.com/forms/d/e/1FAIpQLScPB8emS3Thkp66nvqwmjTEgxp8Y9ufuWTzFyr9kJ5AoI47dQ/viewform?usp=sf_link') + print('='*80) + return "libbitsandbytes_cuda116.dll" # $$$ + + binary_name = "libbitsandbytes_cpu.so" + #if not torch.cuda.is_available(): + #print('No GPU detected. Loading CPU library...') + #return binary_name + + cudart_path = determine_cuda_runtime_lib_path() + if cudart_path is None: + print( + "WARNING: No libcudart.so found! Install CUDA or the cudatoolkit package (anaconda)!" + ) + return binary_name + + print(f"CUDA SETUP: CUDA runtime path found: {cudart_path}") + cuda = get_cuda_lib_handle() + cc = get_compute_capability(cuda) + print(f"CUDA SETUP: Highest compute capability among GPUs detected: {cc}") + cuda_version_string = get_cuda_version(cuda, cudart_path) + + + if cc == '': + print( + "WARNING: No GPU detected! Check your CUDA paths. Processing to load CPU-only library..." + ) + return binary_name + + # 7.5 is the minimum CC vor cublaslt + has_cublaslt = cc in ["7.5", "8.0", "8.6"] + + # TODO: + # (1) CUDA missing cases (no CUDA installed by CUDA driver (nvidia-smi accessible) + # (2) Multiple CUDA versions installed + + # we use ls -l instead of nvcc to determine the cuda version + # since most installations will have the libcudart.so installed, but not the compiler + print(f'CUDA SETUP: Detected CUDA version {cuda_version_string}') + + def get_binary_name(): + "if not has_cublaslt (CC < 7.5), then we have to choose _nocublaslt.so" + bin_base_name = "libbitsandbytes_cuda" + if has_cublaslt: + return f"{bin_base_name}{cuda_version_string}.so" + else: + return f"{bin_base_name}{cuda_version_string}_nocublaslt.so" + + binary_name = get_binary_name() + + return binary_name diff --git a/bitsandbytes_windows/libbitsandbytes_cpu.dll b/bitsandbytes_windows/libbitsandbytes_cpu.dll new file mode 100644 index 00000000..b733af47 Binary files /dev/null and b/bitsandbytes_windows/libbitsandbytes_cpu.dll differ diff --git a/bitsandbytes_windows/libbitsandbytes_cuda116.dll b/bitsandbytes_windows/libbitsandbytes_cuda116.dll new file mode 100644 index 00000000..a999316e Binary files /dev/null and b/bitsandbytes_windows/libbitsandbytes_cuda116.dll differ diff --git a/bitsandbytes_windows/nn/__init__.py b/bitsandbytes_windows/nn/__init__.py new file mode 100644 index 00000000..b1944007 --- /dev/null +++ b/bitsandbytes_windows/nn/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from .modules import Int8Params, Linear8bit, Linear8bitLt +from .modules import Embedding as StableEmbedding diff --git a/codes/models/arch_util.py b/codes/models/arch_util.py index 9383ac7b..4c9caeb0 100644 --- a/codes/models/arch_util.py +++ b/codes/models/arch_util.py @@ -9,6 +9,7 @@ import torch.nn.utils.spectral_norm as SpectralNorm from math import sqrt from utils.util import checkpoint +import torch_intermediary as ml def exists(val): @@ -73,7 +74,7 @@ def initialize_weights(net_l, scale=1): m.weight.data *= scale # for residual block if m.bias is not None: m.bias.data.zero_() - elif isinstance(m, nn.Linear): + elif isinstance(m, ml.Linear): init.kaiming_normal_(m.weight, a=0, mode='fan_in') m.weight.data *= scale if m.bias is not None: @@ -108,7 +109,7 @@ def default_init_weights(module, scale=1): if isinstance(m, nn.Conv2d): kaiming_init(m, a=0, mode='fan_in', bias=0) m.weight.data *= scale - elif isinstance(m, nn.Linear): + elif isinstance(m, ml.Linear): kaiming_init(m, a=0, mode='fan_in', bias=0) m.weight.data *= scale @@ -141,7 +142,7 @@ def linear(*args, **kwargs): """ Create a linear module. """ - return nn.Linear(*args, **kwargs) + return ml.Linear(*args, **kwargs) def avg_pool_nd(dims, *args, **kwargs): diff --git a/codes/models/audio/asr/w2v_wrapper.py b/codes/models/audio/asr/w2v_wrapper.py index 3d578603..f25c562e 100644 --- a/codes/models/audio/asr/w2v_wrapper.py +++ b/codes/models/audio/asr/w2v_wrapper.py @@ -9,6 +9,7 @@ from data.audio.unsupervised_audio_dataset import load_audio from models.audio.tts.tacotron2.text import sequence_to_text from trainer.networks import register_model from utils.util import opt_get +import torch_intermediary as ml def only_letters(string): @@ -51,7 +52,7 @@ class Wav2VecWrapper(nn.Module): self.w2v = Wav2Vec2ForCTC.from_pretrained(basis_model) # Perform some surgery to get the model we actually want. self.w2v.wav2vec2.encoder.gradient_checkpointing = checkpointing_enabled - self.w2v.lm_head = nn.Linear(self.w2v.config.hidden_size, vocab_size) + self.w2v.lm_head = ml.Linear(self.w2v.config.hidden_size, vocab_size) self.w2v.config.vocab_size = vocab_size self.w2v.config.pad_token_id = 0 self.w2v.config.ctc_loss_reduction = 'sum' diff --git a/codes/models/audio/audio_resnet.py b/codes/models/audio/audio_resnet.py index 0d3c32e7..59e8b2fa 100644 --- a/codes/models/audio/audio_resnet.py +++ b/codes/models/audio/audio_resnet.py @@ -5,6 +5,7 @@ import torch.nn as nn from trainer.networks import register_model from utils.util import opt_get from typing import Type, Any, Callable, Union, List, Optional +import torch_intermediary as ml __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', @@ -172,7 +173,7 @@ class ResNet(nn.Module): self.layer4 = self._make_layer(block, 512, layers[3], stride=4, dilate=replace_stride_with_dilation[2]) self.avgpool = nn.AdaptiveAvgPool1d(1) - self.fc = nn.Linear(512 * block.expansion, num_classes) + self.fc = ml.Linear(512 * block.expansion, num_classes) for m in self.modules(): if isinstance(m, nn.Conv1d): diff --git a/codes/models/audio/mel2vec.py b/codes/models/audio/mel2vec.py index 568f0099..25a00a81 100644 --- a/codes/models/audio/mel2vec.py +++ b/codes/models/audio/mel2vec.py @@ -15,13 +15,14 @@ from transformers.deepspeed import is_deepspeed_zero3_enabled from models.arch_util import ResBlock from trainer.networks import register_model from utils.util import checkpoint +import torch_intermediary as ml class Mel2Vec2FeatureProjection(nn.Module): def __init__(self, inner_dim, dropout): super().__init__() self.layer_norm = nn.LayerNorm(inner_dim, eps=1e-5) - self.projection = nn.Linear(inner_dim, inner_dim) + self.projection = ml.Linear(inner_dim, inner_dim) self.dropout = nn.Dropout(dropout) def forward(self, hidden_states): @@ -58,10 +59,10 @@ class Wav2Vec2Attention(nn.Module): self.scaling = self.head_dim**-0.5 self.is_decoder = is_decoder - self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias) - self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias) - self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) - self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.k_proj = ml.Linear(embed_dim, embed_dim, bias=bias) + self.v_proj = ml.Linear(embed_dim, embed_dim, bias=bias) + self.q_proj = ml.Linear(embed_dim, embed_dim, bias=bias) + self.out_proj = ml.Linear(embed_dim, embed_dim, bias=bias) def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() @@ -182,10 +183,10 @@ class Wav2Vec2FeedForward(nn.Module): super().__init__() self.intermediate_dropout = nn.Dropout(dropout) - self.intermediate_dense = nn.Linear(hidden_size, intermediate_size) + self.intermediate_dense = ml.Linear(hidden_size, intermediate_size) self.intermediate_act_fn = F.gelu - self.output_dense = nn.Linear(intermediate_size, hidden_size) + self.output_dense = ml.Linear(intermediate_size, hidden_size) self.output_dropout = nn.Dropout(dropout) def forward(self, hidden_states): @@ -429,7 +430,7 @@ class Mel2Vec(nn.Module): k = math.sqrt(1 / module.projection.in_features) nn.init.uniform_(module.projection.weight, a=-k, b=k) nn.init.uniform_(module.projection.bias, a=-k, b=k) - elif isinstance(module, nn.Linear): + elif isinstance(module, ml.Linear): if self.disable_custom_linear_init: return module.weight.data.normal_(mean=0.0, std=self.linear_init_scale) @@ -510,7 +511,7 @@ class Wav2Vec2GumbelVectorQuantizer(nn.Module): self.codevectors = nn.Parameter( torch.FloatTensor(1, self.num_groups * self.num_vars, codevector_dim // self.num_groups) ) - self.weight_proj = nn.Linear(proj_dim, self.num_groups * self.num_vars) + self.weight_proj = ml.Linear(proj_dim, self.num_groups * self.num_vars) # can be decayed for training self.temperature = 2 @@ -606,8 +607,8 @@ class ContrastiveTrainingWrapper(nn.Module): self.inp_length_factor = inp_length_multiplier # make sure that project_hid & project_q are initialized like normal linear layers - self.project_hid = nn.Linear(inner_dim, self.quantizer.codevector_dim) - self.project_q = nn.Linear(self.quantizer.codevector_dim, self.quantizer.codevector_dim) + self.project_hid = ml.Linear(inner_dim, self.quantizer.codevector_dim) + self.project_q = ml.Linear(self.quantizer.codevector_dim, self.quantizer.codevector_dim) self.reconstruction = do_reconstruction_loss if do_reconstruction_loss: diff --git a/codes/models/audio/music/cheater_gen_ar.py b/codes/models/audio/music/cheater_gen_ar.py index 096e1619..00222217 100644 --- a/codes/models/audio/music/cheater_gen_ar.py +++ b/codes/models/audio/music/cheater_gen_ar.py @@ -2,6 +2,7 @@ import torch import torch.nn.functional as F from torch import nn from transformers import GPT2Config, GPT2Model +import torch_intermediary as ml from models.arch_util import AttentionBlock, ResBlock from models.audio.tts.lucidrains_dvae import DiscreteVAE @@ -55,8 +56,9 @@ class ConditioningAR(nn.Module): self.gpt = GPT2Model(self.config) del self.gpt.wte # Unused, we'll do our own embeddings. - self.embeddings = nn.Embedding(num_vectors, dim) - self.head = nn.Linear(dim, num_vectors) + # nn.Embedding + self.embeddings = ml.Embedding(num_vectors, dim) + self.head = ml.Linear(dim, num_vectors) def forward(self, cheater_codes, conditioning, code_lengths=None, return_latent=False): unused_params = [] diff --git a/codes/models/audio/music/diffwave.py b/codes/models/audio/music/diffwave.py index c7e031c5..71cf765c 100644 --- a/codes/models/audio/music/diffwave.py +++ b/codes/models/audio/music/diffwave.py @@ -17,6 +17,7 @@ import numpy as np import torch import torch.nn as nn import torch.nn.functional as F +import torch_intermediary as ml from math import sqrt @@ -24,7 +25,7 @@ from torch.utils.checkpoint import checkpoint from trainer.networks import register_model -Linear = nn.Linear +Linear = ml.Linear ConvTranspose2d = nn.ConvTranspose2d diff --git a/codes/models/audio/music/flat_diffusion.py b/codes/models/audio/music/flat_diffusion.py index 9def768a..9b8d897b 100644 --- a/codes/models/audio/music/flat_diffusion.py +++ b/codes/models/audio/music/flat_diffusion.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn import torch.nn.functional as F from torch import autocast +import torch_intermediary as ml from models.arch_util import ResBlock from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear @@ -22,7 +23,8 @@ def is_sequence(t): class MultiGroupEmbedding(nn.Module): def __init__(self, tokens, groups, dim): super().__init__() - self.m = nn.ModuleList([nn.Embedding(tokens, dim // groups) for _ in range(groups)]) + # nn.Embedding + self.m = nn.ModuleList([ml.Embedding(tokens, dim // groups) for _ in range(groups)]) def forward(self, x): h = [embedding(x[:, :, i]) for i, embedding in enumerate(self.m)] @@ -158,7 +160,8 @@ class FlatDiffusion(nn.Module): # complex to generate tokens, while generating latents will normally mean propagating through a deep autoregressive # transformer network. if in_groups is None: - self.embeddings = nn.Embedding(token_count, model_channels) + # nn.Embedding + self.embeddings = ml.Embedding(token_count, model_channels) else: self.embeddings = MultiGroupEmbedding(token_count, in_groups, model_channels) self.latent_conditioner = nn.Sequential( diff --git a/codes/models/audio/music/gpt_music.py b/codes/models/audio/music/gpt_music.py index 0e3bfd02..03322df9 100644 --- a/codes/models/audio/music/gpt_music.py +++ b/codes/models/audio/music/gpt_music.py @@ -2,6 +2,7 @@ import torch from torch import nn import torch.nn.functional as F from transformers import GPT2Config, GPT2Model +import torch_intermediary as ml from models.arch_util import AttentionBlock, ResBlock from models.audio.music.music_quantizer import MusicQuantizer @@ -136,8 +137,9 @@ class GptMusicLower(nn.Module): self.gpt = GPT2Model(self.config) del self.gpt.wte # Unused, we'll do our own embeddings. - self.embeddings = nn.ModuleList([nn.Embedding(num_target_vectors, dim // num_vaes) for _ in range(num_vaes)]) - self.heads = nn.ModuleList([nn.Linear(dim, num_target_vectors) for _ in range(num_vaes)]) + # nn.Embedding + self.embeddings = nn.ModuleList([ml.Embedding(num_target_vectors, dim // num_vaes) for _ in range(num_vaes)]) + self.heads = nn.ModuleList([ml.Linear(dim, num_target_vectors) for _ in range(num_vaes)]) def forward(self, mel, conditioning, return_latent=False): unused_params = [] @@ -238,8 +240,9 @@ class GptMusicUpper(nn.Module): self.gpt = GPT2Model(self.config) del self.gpt.wte # Unused, we'll do our own embeddings. - self.embeddings = nn.ModuleList([nn.Embedding(num_upper_vectors, dim // num_upper_groups) for _ in range(num_upper_groups)]) - self.heads = nn.ModuleList([nn.Linear(dim, num_upper_vectors) for _ in range(num_upper_groups)]) + # nn.Embedding + self.embeddings = nn.ModuleList([ml.Embedding(num_upper_vectors, dim // num_upper_groups) for _ in range(num_upper_groups)]) + self.heads = nn.ModuleList([ml.Linear(dim, num_upper_vectors) for _ in range(num_upper_groups)]) def forward(self, mel, conditioning, return_latent=False): diff --git a/codes/models/audio/music/gpt_music2.py b/codes/models/audio/music/gpt_music2.py index acaad51f..d508eb70 100644 --- a/codes/models/audio/music/gpt_music2.py +++ b/codes/models/audio/music/gpt_music2.py @@ -2,6 +2,7 @@ import torch import torch.nn.functional as F from torch import nn from transformers import GPT2Config, GPT2Model +import torch_intermediary as ml from models.arch_util import AttentionBlock, ResBlock from models.audio.tts.lucidrains_dvae import DiscreteVAE @@ -73,8 +74,9 @@ class GptMusicLower(nn.Module): self.gpt = GPT2Model(self.config) del self.gpt.wte # Unused, we'll do our own embeddings. - self.embeddings = nn.ModuleList([nn.Embedding(num_target_vectors, dim // num_vaes) for _ in range(num_vaes)]) - self.heads = nn.ModuleList([nn.Linear(dim, num_target_vectors) for _ in range(num_vaes)]) + # nn.Embedding + self.embeddings = nn.ModuleList([ml.Embedding(num_target_vectors, dim // num_vaes) for _ in range(num_vaes)]) + self.heads = nn.ModuleList([ml.Linear(dim, num_target_vectors) for _ in range(num_vaes)]) def forward(self, mel, return_latent=False): unused_params = [] diff --git a/codes/models/audio/music/instrument_quantizer.py b/codes/models/audio/music/instrument_quantizer.py index 00ea462d..42b5706a 100644 --- a/codes/models/audio/music/instrument_quantizer.py +++ b/codes/models/audio/music/instrument_quantizer.py @@ -3,6 +3,7 @@ import functools import torch import torch.nn as nn import torch.nn.functional as F +import torch_intermediary as ml from models.diffusion.nn import timestep_embedding from models.lucidrains.vq import VectorQuantize @@ -21,8 +22,8 @@ class SelfClassifyingHead(nn.Module): use_rmsnorm=True, ff_glu=True, do_checkpointing=False) self.quantizer = VectorQuantize(out_dim, classes, use_cosine_sim=False, threshold_ema_dead_code=2, sample_codebook_temp=init_temperature) - self.to_output = nn.Linear(dim, out_dim) - self.to_decoder = nn.Linear(out_dim, dim) + self.to_output = ml.Linear(dim, out_dim) + self.to_decoder = ml.Linear(out_dim, dim) def do_ar_step(self, x, used_codes): h = self.dec(x) @@ -90,7 +91,7 @@ class InstrumentQuantizer(nn.Module): """ super().__init__() self.op_dim = op_dim - self.proj = nn.Linear(op_dim, dim) + self.proj = ml.Linear(op_dim, dim) self.encoder = nn.ModuleList([VectorResBlock(dim, dropout) for _ in range(enc_depth)]) self.heads = SelfClassifyingHead(dim, num_classes, op_dim, head_depth, class_seq_len, dropout, max_temp) self.min_gumbel_temperature = min_temp diff --git a/codes/models/audio/music/mel2vec_codes_gpt.py b/codes/models/audio/music/mel2vec_codes_gpt.py index 9cbf9104..62c60768 100644 --- a/codes/models/audio/music/mel2vec_codes_gpt.py +++ b/codes/models/audio/music/mel2vec_codes_gpt.py @@ -1,6 +1,7 @@ import torch from torch import nn import torch.nn.functional as F +import torch_intermediary as ml from transformers import GPT2Config, GPT2Model from trainer.networks import register_model @@ -17,8 +18,9 @@ class Mel2VecCodesGpt(nn.Module): n_inner=dim*2) self.gpt = GPT2Model(self.config) del self.gpt.wte # Unused, we'll do our own embeddings. - self.embeddings = nn.ModuleList([nn.Embedding(num_vectors, dim//num_groups) for _ in range(num_groups)]) - self.heads = nn.ModuleList([nn.Linear(dim, num_vectors) for _ in range(num_groups)]) + # nn.Embedding + self.embeddings = nn.ModuleList([ml.Embedding(num_vectors, dim//num_groups) for _ in range(num_groups)]) + self.heads = nn.ModuleList([ml.Linear(dim, num_vectors) for _ in range(num_groups)]) def forward(self, codes): assert codes.shape[-1] == self.num_groups diff --git a/codes/models/audio/music/music_quantizer.py b/codes/models/audio/music/music_quantizer.py index dd508048..3d3b461b 100644 --- a/codes/models/audio/music/music_quantizer.py +++ b/codes/models/audio/music/music_quantizer.py @@ -3,6 +3,7 @@ import functools import torch from torch import nn import torch.nn.functional as F +import torch_intermediary as ml from models.arch_util import zero_module from models.vqvae.vqvae import Quantize @@ -75,7 +76,7 @@ class Wav2Vec2GumbelVectorQuantizer(nn.Module): self.codevectors = nn.Parameter( torch.FloatTensor(1, self.num_groups * self.num_vars, codevector_dim // self.num_groups) ) - self.weight_proj = nn.Linear(proj_dim, self.num_groups * self.num_vars) + self.weight_proj = ml.Linear(proj_dim, self.num_groups * self.num_vars) # can be decayed for training self.temperature = 2 diff --git a/codes/models/audio/music/music_quantizer2.py b/codes/models/audio/music/music_quantizer2.py index 8fa73c65..d7df3658 100644 --- a/codes/models/audio/music/music_quantizer2.py +++ b/codes/models/audio/music/music_quantizer2.py @@ -3,6 +3,7 @@ import functools import torch from torch import nn import torch.nn.functional as F +import torch_intermediary as ml from models.arch_util import zero_module from models.vqvae.vqvae import Quantize @@ -87,7 +88,7 @@ class Wav2Vec2GumbelVectorQuantizer(nn.Module): self.codevectors = nn.Parameter( torch.FloatTensor(1, self.num_groups * self.num_vars, codevector_dim // self.num_groups) ) - self.weight_proj = nn.Linear(proj_dim, self.num_groups * self.num_vars) + self.weight_proj = ml.Linear(proj_dim, self.num_groups * self.num_vars) # can be decayed for training self.temperature = 2 diff --git a/codes/models/audio/music/tfdpc_v5.py b/codes/models/audio/music/tfdpc_v5.py index a71650ae..998906df 100644 --- a/codes/models/audio/music/tfdpc_v5.py +++ b/codes/models/audio/music/tfdpc_v5.py @@ -1,3 +1,4 @@ + import itertools import os import random @@ -7,6 +8,7 @@ import torch.nn as nn import torch.nn.functional as F import torchaudio import torchvision +import torch_intermediary as ml from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear from models.diffusion.unet_diffusion import TimestepBlock @@ -54,12 +56,12 @@ class ConcatAttentionBlock(TimestepBlock): self.prenorm = RMSScaleShiftNorm(trunk_dim, embed_dim=time_embed_dim, bias=False) if cond_projection: self.tdim = trunk_dim+cond_dim_hidden - self.cond_project = nn.Linear(cond_dim_in, cond_dim_hidden) + self.cond_project = ml.Linear(cond_dim_in, cond_dim_hidden) else: self.tdim = trunk_dim self.block1 = SubBlock(self.tdim, contraction_dim, heads, dropout, use_conv) self.block2 = SubBlock(self.tdim+contraction_dim*2, contraction_dim, heads, dropout, use_conv) - self.out = nn.Linear(contraction_dim*4, trunk_dim, bias=False) + self.out = ml.Linear(contraction_dim*4, trunk_dim, bias=False) self.out.weight.data.zero_() def forward(self, x, cond, timestep_emb, rotary_emb): @@ -87,7 +89,7 @@ class ConditioningEncoder(nn.Module): self.init = nn.Conv1d(cond_dim, embedding_dim, kernel_size=1) self.time_proj = time_proj if time_proj: - self.time_proj = nn.Linear(time_embed_dim, embedding_dim) + self.time_proj = ml.Linear(time_embed_dim, embedding_dim) self.attn = Encoder( dim=embedding_dim, depth=attn_blocks, diff --git a/codes/models/audio/music/transformer_diffusion12.py b/codes/models/audio/music/transformer_diffusion12.py index 056e6426..ff3ccc90 100644 --- a/codes/models/audio/music/transformer_diffusion12.py +++ b/codes/models/audio/music/transformer_diffusion12.py @@ -4,6 +4,7 @@ from time import time import torch import torch.nn as nn import torch.nn.functional as F +import torch_intermediary as ml from models.arch_util import ResBlock from models.audio.music.gpt_music2 import UpperEncoder, GptMusicLower @@ -27,7 +28,8 @@ def is_sequence(t): class MultiGroupEmbedding(nn.Module): def __init__(self, tokens, groups, dim): super().__init__() - self.m = nn.ModuleList([nn.Embedding(tokens, dim // groups) for _ in range(groups)]) + # nn.Embedding + self.m = nn.ModuleList([ml.Embedding(tokens, dim // groups) for _ in range(groups)]) def forward(self, x): h = [embedding(x[:, :, i]) for i, embedding in enumerate(self.m)] @@ -68,7 +70,7 @@ class ConcatAttentionBlock(TimestepBlock): self.prenorm = RMSScaleShiftNorm(trunk_dim, embed_dim=time_embed_dim, bias=False) self.block1 = SubBlock(trunk_dim, contraction_dim, heads, dropout) self.block2 = SubBlock(trunk_dim+contraction_dim*2, contraction_dim, heads, dropout) - self.out = nn.Linear(contraction_dim*4, trunk_dim, bias=False) + self.out = ml.Linear(contraction_dim*4, trunk_dim, bias=False) self.out.weight.data.zero_() def forward(self, x, timestep_emb, rotary_emb): @@ -129,7 +131,7 @@ class TransformerDiffusion(nn.Module): ) prenet_heads = prenet_channels//64 - self.input_converter = nn.Linear(input_vec_dim, prenet_channels) + self.input_converter = ml.Linear(input_vec_dim, prenet_channels) self.code_converter = Encoder( dim=prenet_channels, depth=prenet_layers, @@ -145,7 +147,7 @@ class TransformerDiffusion(nn.Module): self.unconditioned_embedding = nn.Parameter(torch.randn(1,1,prenet_channels)) self.rotary_embeddings = RotaryEmbedding(rotary_emb_dim) - self.intg = nn.Linear(prenet_channels*2, model_channels) + self.intg = ml.Linear(prenet_channels*2, model_channels) self.layers = TimestepRotaryEmbedSequential(*[ConcatAttentionBlock(model_channels, contraction_dim, time_embed_dim, num_heads, dropout) for _ in range(num_layers)]) self.out = nn.Sequential( diff --git a/codes/models/audio/music/transformer_diffusion13.py b/codes/models/audio/music/transformer_diffusion13.py index e641b3e1..c946af4e 100644 --- a/codes/models/audio/music/transformer_diffusion13.py +++ b/codes/models/audio/music/transformer_diffusion13.py @@ -5,6 +5,7 @@ from random import randrange import torch import torch.nn as nn import torch.nn.functional as F +import torch_intermediary as ml from models.arch_util import ResBlock, TimestepEmbedSequential, AttentionBlock, build_local_attention_mask, cGLU, \ RelativeQKBias @@ -69,13 +70,14 @@ class ConditioningEncoder(nn.Module): super().__init__() attn = [] self.init = nn.Conv1d(spec_dim, hidden_dim, kernel_size=5, stride=2) - self.resolution_embedding = nn.Embedding(num_resolutions, hidden_dim) + # nn.Embedding + self.resolution_embedding = ml.Embedding(num_resolutions, hidden_dim) self.resolution_embedding.weight.data.mul(.1) # Reduces the relative influence of this embedding from the start. for a in range(attn_blocks): attn.append(AttentionBlock(hidden_dim, num_attn_heads, do_checkpoint=do_checkpointing)) attn.append(ResBlock(hidden_dim, dims=1, checkpointing_enabled=do_checkpointing)) self.attn = nn.Sequential(*attn) - self.out = nn.Linear(hidden_dim, out_dim, bias=False) + self.out = ml.Linear(hidden_dim, out_dim, bias=False) self.dim = hidden_dim self.do_checkpointing = do_checkpointing @@ -131,7 +133,8 @@ class TransformerDiffusion(nn.Module): nn.SiLU(), linear(time_embed_dim, time_proj_dim), ) - self.resolution_embed = nn.Embedding(resolution_steps, time_proj_dim) + # nn.Embedding + self.resolution_embed = ml.Embedding(resolution_steps, time_proj_dim) self.conditioning_encoder = ConditioningEncoder(in_channels, model_channels, cond_proj_dim, resolution_steps, num_attn_heads=model_channels//64) self.unconditioned_embedding = nn.Parameter(torch.randn(1,cond_proj_dim)) diff --git a/codes/models/audio/music/unet_diffusion_music_codes.py b/codes/models/audio/music/unet_diffusion_music_codes.py index d0818744..5a135502 100644 --- a/codes/models/audio/music/unet_diffusion_music_codes.py +++ b/codes/models/audio/music/unet_diffusion_music_codes.py @@ -8,6 +8,7 @@ import torch as th import torch.nn as nn import torch.nn.functional as F import torchvision # For debugging, not actually used. +import torch_intermediary as ml from models.audio.music.gpt_music import GptMusicLower from models.audio.music.music_quantizer import MusicQuantizer @@ -490,7 +491,7 @@ class UNetMusicModel(nn.Module): ) if self.ar_prior: - self.ar_input = nn.Linear(input_vec_dim, model_channels) + self.ar_input = ml.Linear(input_vec_dim, model_channels) self.ar_prior_intg = Encoder( dim=model_channels, depth=4, @@ -504,7 +505,7 @@ class UNetMusicModel(nn.Module): ff_mult=1, ) else: - self.input_converter = nn.Linear(input_vec_dim, model_channels) + self.input_converter = ml.Linear(input_vec_dim, model_channels) self.code_converter = Encoder( dim=model_channels, depth=4, @@ -521,7 +522,8 @@ class UNetMusicModel(nn.Module): self.x_processor = conv_nd(dims, in_channels, model_channels, 3, padding=1) if self.num_classes is not None: - self.label_emb = nn.Embedding(num_classes, time_embed_dim) + # nn.Embedding + self.label_emb = ml.Embedding(num_classes, time_embed_dim) self.use_raw_y_as_embedding = use_raw_y_as_embedding assert not ((self.num_classes is not None) and use_raw_y_as_embedding) # These are mutually-exclusive. diff --git a/codes/models/audio/tts/ctc_code_generator.py b/codes/models/audio/tts/ctc_code_generator.py index 68905115..801c3fa2 100644 --- a/codes/models/audio/tts/ctc_code_generator.py +++ b/codes/models/audio/tts/ctc_code_generator.py @@ -3,6 +3,7 @@ from random import random import torch import torch.nn as nn import torch.nn.functional as F +import torch_intermediary as ml from models.audio.tts.unet_diffusion_tts7 import CheckpointedLayer from models.lucidrains.x_transformers import Encoder @@ -36,9 +37,12 @@ class CtcCodeGenerator(nn.Module): self.ctc_codes = ctc_codes pred_codes = (max_pad+1)*(max_repeat+1) - self.position_embedding = nn.Embedding(max_length, model_dim) - self.codes_embedding = nn.Embedding(ctc_codes, model_dim) - self.recursive_embedding = nn.Embedding(pred_codes, model_dim) + # nn.Embedding + self.position_embedding = ml.Embedding(max_length, model_dim) + # nn.Embedding + self.codes_embedding = ml.Embedding(ctc_codes, model_dim) + # nn.Embedding + self.recursive_embedding = ml.Embedding(pred_codes, model_dim) self.mask_embedding = nn.Parameter(torch.randn(model_dim)) self.encoder = Encoder( dim=model_dim, @@ -50,8 +54,8 @@ class CtcCodeGenerator(nn.Module): ff_glu=True, rotary_pos_emb=True, ) - self.pred_head = nn.Linear(model_dim, pred_codes) - self.confidence_head = nn.Linear(model_dim, 1) + self.pred_head = ml.Linear(model_dim, pred_codes) + self.confidence_head = ml.Linear(model_dim, 1) def inference(self, codes, pads, repeats): position_h = self.position_embedding(torch.arange(0, codes.shape[-1], device=codes.device)) diff --git a/codes/models/audio/tts/diffusion_encoder.py b/codes/models/audio/tts/diffusion_encoder.py index 1dc91c0d..15c3a2f2 100644 --- a/codes/models/audio/tts/diffusion_encoder.py +++ b/codes/models/audio/tts/diffusion_encoder.py @@ -5,6 +5,8 @@ from functools import partial import torch import torch.nn as nn +import torch_intermediary as ml + from x_transformers.x_transformers import groupby_prefix_and_trim, FixedPositionalEmbedding, default, RotaryEmbedding, \ DEFAULT_DIM_HEAD, RelativePositionBias, LearnedAlibiPositionalBias, AlibiPositionalBias, ScaleNorm, RMSNorm, \ exists, Attention, FeedForward, Scale, ShiftTokens, GRUGating, Residual, cast_tuple, equals, LayerIntermediates, \ @@ -16,7 +18,7 @@ class TimeIntegrationBlock(nn.Module): super().__init__() self.emb_layers = nn.Sequential( nn.SiLU(), - nn.Linear( + ml.Linear( time_emb_dim, 2 * dim ), diff --git a/codes/models/audio/tts/mini_encoder.py b/codes/models/audio/tts/mini_encoder.py index 23283173..4a61199b 100644 --- a/codes/models/audio/tts/mini_encoder.py +++ b/codes/models/audio/tts/mini_encoder.py @@ -1,5 +1,6 @@ import torch import torch.nn as nn +import torch_intermediary as ml from models.diffusion.nn import normalization, conv_nd, zero_module @@ -138,7 +139,7 @@ class AudioMiniEncoderWithClassifierHead(nn.Module): def __init__(self, classes, distribute_zero_label=True, **kwargs): super().__init__() self.enc = AudioMiniEncoder(**kwargs) - self.head = nn.Linear(self.enc.dim, classes) + self.head = ml.Linear(self.enc.dim, classes) self.num_classes = classes self.distribute_zero_label = distribute_zero_label @@ -183,7 +184,7 @@ class QueryProvidedAttentionBlock(nn.Module): ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}" self.num_heads = channels // num_head_channels self.norm = normalization(channels) - self.q = nn.Linear(channels, channels) + self.q = ml.Linear(channels, channels) self.qnorm = nn.LayerNorm(channels) self.kv = conv_nd(1, channels, channels*2, 1) if use_new_attention_order: diff --git a/codes/models/audio/tts/random_latent_converter.py b/codes/models/audio/tts/random_latent_converter.py index d4b5dd00..9f14ebb3 100644 --- a/codes/models/audio/tts/random_latent_converter.py +++ b/codes/models/audio/tts/random_latent_converter.py @@ -3,6 +3,7 @@ import math import torch import torch.nn as nn import torch.nn.functional as F +import torch_intermediary as ml from trainer.networks import register_model from utils.util import opt_get @@ -44,7 +45,7 @@ class RandomLatentConverter(nn.Module): def __init__(self, channels): super().__init__() self.layers = nn.Sequential(*[EqualLinear(channels, channels, lr_mul=.1) for _ in range(5)], - nn.Linear(channels, channels)) + ml.Linear(channels, channels)) self.channels = channels def forward(self, ref): diff --git a/codes/models/audio/tts/tacotron2/layers.py b/codes/models/audio/tts/tacotron2/layers.py index 8b69e0df..11022c02 100644 --- a/codes/models/audio/tts/tacotron2/layers.py +++ b/codes/models/audio/tts/tacotron2/layers.py @@ -3,12 +3,13 @@ from librosa.filters import mel as librosa_mel_fn from models.audio.tts.tacotron2.audio_processing import dynamic_range_compression from models.audio.tts.tacotron2.audio_processing import dynamic_range_decompression from models.audio.tts.tacotron2.stft import STFT +import torch_intermediary as ml class LinearNorm(torch.nn.Module): def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'): super(LinearNorm, self).__init__() - self.linear_layer = torch.nn.Linear(in_dim, out_dim, bias=bias) + self.linear_layer = torch.ml.Linear(in_dim, out_dim, bias=bias) torch.nn.init.xavier_uniform_( self.linear_layer.weight, diff --git a/codes/models/audio/tts/tacotron2/tacotron2.py b/codes/models/audio/tts/tacotron2/tacotron2.py index 73ecc754..f8f64cb7 100644 --- a/codes/models/audio/tts/tacotron2/tacotron2.py +++ b/codes/models/audio/tts/tacotron2/tacotron2.py @@ -8,6 +8,7 @@ from models.audio.tts.tacotron2.layers import ConvNorm, LinearNorm from models.audio.tts.tacotron2.hparams import create_hparams from trainer.networks import register_model from models.audio.tts.tacotron2.taco_utils import get_mask_from_lengths +import torch_intermediary as ml class LocationLayer(nn.Module): @@ -463,7 +464,8 @@ class Tacotron2(nn.Module): self.fp16_run = hparams.fp16_run self.n_mel_channels = hparams.n_mel_channels self.n_frames_per_step = hparams.n_frames_per_step - self.embedding = nn.Embedding( + # nn.Embedding + self.embedding = ml.Embedding( hparams.n_symbols, hparams.symbols_embedding_dim) std = sqrt(2.0 / (hparams.n_symbols + hparams.symbols_embedding_dim)) val = sqrt(3.0) * std # uniform bounds for std diff --git a/codes/models/audio/tts/tacotron2/wave_tacotron.py b/codes/models/audio/tts/tacotron2/wave_tacotron.py index 4c02e05e..8e73e93d 100644 --- a/codes/models/audio/tts/tacotron2/wave_tacotron.py +++ b/codes/models/audio/tts/tacotron2/wave_tacotron.py @@ -13,6 +13,7 @@ from models.audio.tts.tacotron2.tacotron2 import Attention, Encoder from trainer.networks import register_model from models.audio.tts.tacotron2.taco_utils import get_mask_from_lengths from utils.util import checkpoint +import torch_intermediary as ml @@ -185,7 +186,8 @@ class WaveTacotron2(nn.Module): self.fp16_run = hparams.fp16_run self.n_mel_channels = hparams.n_mel_channels self.n_frames_per_step = hparams.n_frames_per_step - self.embedding = nn.Embedding( + # nn.Embedding + self.embedding = ml.Embedding( hparams.n_symbols, hparams.symbols_embedding_dim) std = sqrt(2.0 / (hparams.n_symbols + hparams.symbols_embedding_dim)) val = sqrt(3.0) * std # uniform bounds for std diff --git a/codes/models/audio/tts/transformer_builders.py b/codes/models/audio/tts/transformer_builders.py index 238dab66..ce88d1a7 100644 --- a/codes/models/audio/tts/transformer_builders.py +++ b/codes/models/audio/tts/transformer_builders.py @@ -25,6 +25,7 @@ import random from time import time import torch import torch.nn as nn +import torch_intermediary as ml from tqdm import tqdm @@ -35,7 +36,8 @@ def null_position_embeddings(range, dim): class LearnedPositionEmbeddings(nn.Module): def __init__(self, seq_len, model_dim, init=.02, relative=False): super().__init__() - self.emb = nn.Embedding(seq_len, model_dim) + # nn.Embedding + self.emb = ml.Embedding(seq_len, model_dim) # Initializing this way is standard for GPT-2 self.emb.weight.data.normal_(mean=0.0, std=init) self.relative = relative diff --git a/codes/models/audio/tts/transformer_diffusion_tts.py b/codes/models/audio/tts/transformer_diffusion_tts.py index 430ebc4a..bb78a008 100644 --- a/codes/models/audio/tts/transformer_diffusion_tts.py +++ b/codes/models/audio/tts/transformer_diffusion_tts.py @@ -7,6 +7,7 @@ from models.diffusion.unet_diffusion import TimestepEmbedSequential, TimestepBlo from models.lucidrains.x_transformers import Encoder, Attention, FeedForward, RMSScaleShiftNorm, RotaryEmbedding from trainer.networks import register_model from utils.util import checkpoint +import torch_intermediary as ml def is_latent(t): @@ -19,7 +20,8 @@ def is_sequence(t): class MultiGroupEmbedding(nn.Module): def __init__(self, tokens, groups, dim): super().__init__() - self.m = nn.ModuleList([nn.Embedding(tokens, dim // groups) for _ in range(groups)]) + # nn.Embedding + self.m = nn.ModuleList([ml.Embedding(tokens, dim // groups) for _ in range(groups)]) def forward(self, x): h = [embedding(x[:, :, i]) for i, embedding in enumerate(self.m)] @@ -100,15 +102,17 @@ class TransformerDiffusionTTS(nn.Module): ff_glu=True, rotary_pos_emb=True, ) - self.clvp_encoder = nn.Linear(clvp_in_dim, model_channels) - self.type_embedding = nn.Embedding(types, model_channels) + self.clvp_encoder = ml.Linear(clvp_in_dim, model_channels) + # nn.Embedding + self.type_embedding = ml.Embedding(types, model_channels) # Either code_converter or latent_converter is used, depending on what type of conditioning data is fed. # This model is meant to be able to be trained on both for efficiency purposes - it is far less computationally # complex to generate tokens, while generating latents will normally mean propagating through a deep autoregressive # transformer network. if in_groups is None: - self.embeddings = nn.Embedding(token_count, model_channels) + # nn.Embedding + self.embeddings = ml.Embedding(token_count, model_channels) else: self.embeddings = MultiGroupEmbedding(token_count, in_groups, model_channels) self.latent_conditioner = nn.Sequential( @@ -140,7 +144,7 @@ class TransformerDiffusionTTS(nn.Module): self.mel_head = nn.Conv1d(model_channels, in_channels, kernel_size=3, padding=1) self.rotary_embeddings = RotaryEmbedding(rotary_emb_dim) - self.intg = nn.Linear(model_channels*2, model_channels) + self.intg = ml.Linear(model_channels*2, model_channels) self.layers = TimestepRotaryEmbedSequential(*[AttentionBlock(model_channels, model_channels//64, dropout) for _ in range(num_layers)]) self.out = nn.Sequential( diff --git a/codes/models/audio/tts/transformer_diffusion_tts2.py b/codes/models/audio/tts/transformer_diffusion_tts2.py index a1351539..e4cb3a6e 100644 --- a/codes/models/audio/tts/transformer_diffusion_tts2.py +++ b/codes/models/audio/tts/transformer_diffusion_tts2.py @@ -1,6 +1,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +import torch_intermediary as ml from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear from models.diffusion.unet_diffusion import TimestepEmbedSequential, TimestepBlock @@ -19,7 +20,8 @@ def is_sequence(t): class MultiGroupEmbedding(nn.Module): def __init__(self, tokens, groups, dim): super().__init__() - self.m = nn.ModuleList([nn.Embedding(tokens, dim // groups) for _ in range(groups)]) + # nn.Embedding + self.m = nn.ModuleList([ml.Embedding(tokens, dim // groups) for _ in range(groups)]) def forward(self, x): h = [embedding(x[:, :, i]) for i, embedding in enumerate(self.m)] @@ -40,7 +42,7 @@ class DietAttentionBlock(TimestepBlock): def __init__(self, in_dim, dim, heads, dropout): super().__init__() self.rms_scale_norm = RMSScaleShiftNorm(in_dim) - self.proj = nn.Linear(in_dim, dim) + self.proj = ml.Linear(in_dim, dim) self.attn = Attention(dim, heads=heads, causal=False, dropout=dropout) self.ff = FeedForward(dim, in_dim, mult=1, dropout=dropout, zero_init_output=True) @@ -105,15 +107,17 @@ class TransformerDiffusionTTS(nn.Module): ff_glu=True, rotary_pos_emb=True, ) - self.clvp_encoder = nn.Linear(clvp_in_dim, prenet_channels) - self.type_embedding = nn.Embedding(types, prenet_channels) + self.clvp_encoder = ml.Linear(clvp_in_dim, prenet_channels) + # nn.Embedding + self.type_embedding = ml.Embedding(types, prenet_channels) # Either code_converter or latent_converter is used, depending on what type of conditioning data is fed. # This model is meant to be able to be trained on both for efficiency purposes - it is far less computationally # complex to generate tokens, while generating latents will normally mean propagating through a deep autoregressive # transformer network. if in_groups is None: - self.embeddings = nn.Embedding(token_count, prenet_channels) + # nn.Embedding + self.embeddings = ml.Embedding(token_count, prenet_channels) else: self.embeddings = MultiGroupEmbedding(token_count, in_groups, prenet_channels) self.latent_conditioner = nn.Sequential( @@ -144,8 +148,8 @@ class TransformerDiffusionTTS(nn.Module): self.unconditioned_embedding = nn.Parameter(torch.randn(1,1,prenet_channels)) self.rotary_embeddings = RotaryEmbedding(rotary_emb_dim) - self.cond_intg = nn.Linear(prenet_channels*4, model_channels) - self.intg = nn.Linear(prenet_channels*2, model_channels) + self.cond_intg = ml.Linear(prenet_channels*4, model_channels) + self.intg = ml.Linear(prenet_channels*2, model_channels) self.layers = TimestepRotaryEmbedSequential(*[DietAttentionBlock(model_channels, block_channels, block_channels // 64, dropout) for _ in range(num_layers)]) diff --git a/codes/models/audio/tts/unet_diffusion_tts7.py b/codes/models/audio/tts/unet_diffusion_tts7.py index a323c9b0..02d25cb3 100644 --- a/codes/models/audio/tts/unet_diffusion_tts7.py +++ b/codes/models/audio/tts/unet_diffusion_tts7.py @@ -5,6 +5,7 @@ import torch import torch.nn as nn import torch.nn.functional as F from torch import autocast +import torch_intermediary as ml from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear from models.diffusion.unet_diffusion import AttentionBlock, TimestepEmbedSequential, \ @@ -247,14 +248,16 @@ class DiffusionTts(nn.Module): ) embedding_dim = model_channels * 8 - self.code_embedding = nn.Embedding(num_tokens+1, embedding_dim) + # nn.Embedding + self.code_embedding = ml.Embedding(num_tokens+1, embedding_dim) self.contextual_embedder = AudioMiniEncoder(1, embedding_dim, base_channels=32, depth=6, resnet_blocks=1, attn_blocks=2, num_attn_heads=2, dropout=dropout, downsample_factor=4, kernel_size=5) self.conditioning_conv = nn.Conv1d(embedding_dim*3, embedding_dim, 1) self.enable_unaligned_inputs = enabled_unaligned_inputs if enabled_unaligned_inputs: - self.unaligned_embedder = nn.Embedding(num_unaligned_tokens, embedding_dim) + # nn.Embedding + self.unaligned_embedder = ml.Embedding(num_unaligned_tokens, embedding_dim) self.unaligned_encoder = CheckpointedXTransformerEncoder( max_seq_len=-1, use_pos_emb=False, diff --git a/codes/models/audio/tts/unet_diffusion_tts9.py b/codes/models/audio/tts/unet_diffusion_tts9.py index 278fde70..a00b2758 100644 --- a/codes/models/audio/tts/unet_diffusion_tts9.py +++ b/codes/models/audio/tts/unet_diffusion_tts9.py @@ -5,6 +5,7 @@ import torch.nn as nn import torch.nn.functional as F from torch import autocast from x_transformers import Encoder +import torch_intermediary as ml from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear from models.diffusion.unet_diffusion import AttentionBlock, TimestepEmbedSequential, \ @@ -206,7 +207,8 @@ class DiffusionTts(nn.Module): # complex to generate tokens, while generating latents will normally mean propagating through a deep autoregressive # transformer network. self.code_converter = nn.Sequential( - nn.Embedding(in_tokens, conditioning_dim), + # nn.Embedding + ml.Embedding(in_tokens, conditioning_dim), CheckpointedXTransformerEncoder( needs_permute=False, max_seq_len=-1, diff --git a/codes/models/audio/tts/unet_diffusion_tts_flat.py b/codes/models/audio/tts/unet_diffusion_tts_flat.py index 6b034680..34017011 100644 --- a/codes/models/audio/tts/unet_diffusion_tts_flat.py +++ b/codes/models/audio/tts/unet_diffusion_tts_flat.py @@ -5,6 +5,7 @@ import torch import torch.nn as nn import torch.nn.functional as F from torch import autocast +import torch_intermediary as ml from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear from models.diffusion.unet_diffusion import TimestepEmbedSequential, TimestepBlock, QKVAttentionLegacy @@ -193,7 +194,9 @@ class DiffusionTtsFlat(nn.Module): # This model is meant to be able to be trained on both for efficiency purposes - it is far less computationally # complex to generate tokens, while generating latents will normally mean propagating through a deep autoregressive # transformer network. - self.code_embedding = nn.Embedding(in_tokens, model_channels) + + # nn.Embedding + self.code_embedding = ml.Embedding(in_tokens, model_channels) self.code_converter = nn.Sequential( AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), diff --git a/codes/models/audio/tts/unified_voice2.py b/codes/models/audio/tts/unified_voice2.py index 435526c6..80e7a394 100644 --- a/codes/models/audio/tts/unified_voice2.py +++ b/codes/models/audio/tts/unified_voice2.py @@ -1,6 +1,7 @@ import torch import torch.nn as nn import torch.nn.functional as F + from transformers import GPT2Config, GPT2PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions from transformers.models.gpt2.modeling_gpt2 import GPT2Attention @@ -12,6 +13,7 @@ from models.lucidrains.x_transformers import RotaryEmbedding, apply_rotary_pos_e from trainer.networks import register_model from utils.util import opt_get +import torch_intermediary as ml class ResBlock(nn.Module): """ @@ -279,9 +281,11 @@ class UnifiedVoice(nn.Module): self.mel_length_compression = mel_length_compression self.conditioning_encoder = ConditioningEncoder(80, model_dim, num_attn_heads=heads) self.average_conditioning_embeddings = average_conditioning_embeddings - self.text_embedding = nn.Embedding(self.number_text_tokens, model_dim) + # nn.Embedding + self.text_embedding = ml.Embedding(self.number_text_tokens, model_dim) if use_mel_codes_as_input: - self.mel_embedding = nn.Embedding(self.number_mel_codes, model_dim) + # nn.Embedding + self.mel_embedding = ml.Embedding(self.number_mel_codes, model_dim) else: self.mel_embedding = MelEncoder(model_dim, resblocks_per_reduction=1) self.gpt, self.mel_pos_embedding, self.text_pos_embedding, self.mel_layer_pos_embedding, self.text_layer_pos_embedding = \ @@ -294,8 +298,8 @@ class UnifiedVoice(nn.Module): self.text_solo_embedding = 0 self.final_norm = nn.LayerNorm(model_dim) - self.text_head = nn.Linear(model_dim, self.number_text_tokens) - self.mel_head = nn.Linear(model_dim, self.number_mel_codes) + self.text_head = ml.Linear(model_dim, self.number_text_tokens) + self.mel_head = ml.Linear(model_dim, self.number_mel_codes) # Initialize the embeddings per the GPT-2 scheme embeddings = [self.text_embedding] diff --git a/codes/models/audio/tts/unified_voice3.py b/codes/models/audio/tts/unified_voice3.py index dd26c789..49e2258c 100644 --- a/codes/models/audio/tts/unified_voice3.py +++ b/codes/models/audio/tts/unified_voice3.py @@ -1,6 +1,8 @@ import torch import torch.nn as nn import torch.nn.functional as F +import torch_intermediary as ml + from transformers import GPT2Config, GPT2PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions from transformers.models.gpt2.modeling_gpt2 import GPT2Attention @@ -271,15 +273,17 @@ class UnifiedVoice(nn.Module): self.model_dim = model_dim self.mel_length_compression = mel_length_compression self.conditioning_encoder = ConditioningEncoder(80, model_dim, num_attn_heads=heads) - self.text_embedding = nn.Embedding(self.number_text_tokens*types+1, model_dim) - self.mel_embedding = nn.Embedding(self.number_mel_codes, model_dim) + # nn.Embedding + self.text_embedding = ml.Embedding(self.number_text_tokens*types+1, model_dim) + # nn.Embedding + self.mel_embedding = ml.Embedding(self.number_mel_codes, model_dim) self.gpt, self.mel_pos_embedding, self.text_pos_embedding, self.mel_layer_pos_embedding, self.text_layer_pos_embedding = \ build_hf_gpt_transformer(layers, model_dim, heads, self.max_mel_tokens, self.max_text_tokens, checkpointing) self.final_norm = nn.LayerNorm(model_dim) - self.text_head = nn.Linear(model_dim, self.number_text_tokens*types+1) - self.mel_head = nn.Linear(model_dim, self.number_mel_codes) - self.aligned_head = nn.Linear(model_dim, number_aligned_text_codes) + self.text_head = ml.Linear(model_dim, self.number_text_tokens*types+1) + self.mel_head = ml.Linear(model_dim, self.number_mel_codes) + self.aligned_head = ml.Linear(model_dim, number_aligned_text_codes) # Initialize the embeddings per the GPT-2 scheme embeddings = [self.text_embedding, self.mel_embedding] diff --git a/codes/models/audio/tts/unified_voice4.py b/codes/models/audio/tts/unified_voice4.py index a186260d..9d8a8568 100644 --- a/codes/models/audio/tts/unified_voice4.py +++ b/codes/models/audio/tts/unified_voice4.py @@ -11,6 +11,7 @@ from models.audio.tts.transformer_builders import build_hf_gpt_transformer from models.lucidrains.x_transformers import RotaryEmbedding, apply_rotary_pos_emb from trainer.networks import register_model from utils.util import opt_get +import torch_intermediary as ml class ResBlock(nn.Module): @@ -255,15 +256,17 @@ class UnifiedVoice(nn.Module): self.model_dim = model_dim self.mel_length_compression = mel_length_compression self.conditioning_encoder = ConditioningEncoder(80, model_dim, num_attn_heads=heads) - self.text_embedding = nn.Embedding(self.number_text_tokens*types+1, model_dim) - self.mel_embedding = nn.Embedding(self.number_mel_codes, model_dim) + # nn.Embedding + self.text_embedding = ml.Embedding(self.number_text_tokens*types+1, model_dim) + # nn.Embedding + self.mel_embedding = ml.Embedding(self.number_mel_codes, model_dim) self.gpt, self.mel_pos_embedding, self.text_pos_embedding, self.mel_layer_pos_embedding, self.text_layer_pos_embedding = \ build_hf_gpt_transformer(layers, model_dim, heads, self.max_mel_tokens, self.max_text_tokens, checkpointing) self.final_norm = nn.LayerNorm(model_dim) - self.text_head = nn.Linear(model_dim, self.number_text_tokens*types+1) - self.mel_head = nn.Linear(model_dim, self.number_mel_codes) - self.alignment_head = nn.Linear(model_dim, 256) + self.text_head = ml.Linear(model_dim, self.number_text_tokens*types+1) + self.mel_head = ml.Linear(model_dim, self.number_mel_codes) + self.alignment_head = ml.Linear(model_dim, 256) if only_alignment_head: for p in self.parameters(): diff --git a/codes/models/audio/tts/voice_voice_clip.py b/codes/models/audio/tts/voice_voice_clip.py index 22fa2d80..842f8b17 100644 --- a/codes/models/audio/tts/voice_voice_clip.py +++ b/codes/models/audio/tts/voice_voice_clip.py @@ -8,6 +8,7 @@ from models.audio.tts.mini_encoder import AudioMiniEncoder from trainer.injectors.spec_augment import spec_augment from trainer.networks import register_model from utils.util import opt_get +import torch_intermediary as ml def exists(val): @@ -36,7 +37,7 @@ class VoiceCLIP(nn.Module): self.encoder = AudioMiniEncoder(80, encoder_output) if pretrained_encoder_dict_path is not None: self.encoder.load_state_dict(torch.load(pretrained_encoder_dict_path)) - self.to_latent = nn.Linear(encoder_output, dim_latent, bias=False) + self.to_latent = ml.Linear(encoder_output, dim_latent, bias=False) self.temperature = nn.Parameter(torch.tensor(1.)) self.mel_compression_ratio = mel_compression_ratio diff --git a/codes/models/audio/tts/w2v_matcher.py b/codes/models/audio/tts/w2v_matcher.py index a2d261a6..636a1ca0 100644 --- a/codes/models/audio/tts/w2v_matcher.py +++ b/codes/models/audio/tts/w2v_matcher.py @@ -7,6 +7,7 @@ from x_transformers import Encoder, Decoder, ContinuousTransformerWrapper from models.audio.tts.mini_encoder import AudioMiniEncoder from trainer.networks import register_model +import torch_intermediary as ml class CheckpointedLayer(nn.Module): @@ -56,7 +57,8 @@ class Wav2VecMatcher(nn.Module): WAV2VEC_CHANNELS = 1024 self.conditioning_encoder = AudioMiniEncoder(1, model_dim, base_channels=32, depth=6, resnet_blocks=1, attn_blocks=2, num_attn_heads=2, dropout=dropout, downsample_factor=4, kernel_size=5) - self.text_embedding = nn.Embedding(num_text_tokens, model_dim) + # nn.Embedding + self.text_embedding = ml.Embedding(num_text_tokens, model_dim) self.encoder = CheckpointedXTransformer( max_seq_len=-1, use_pos_emb=False, @@ -73,8 +75,8 @@ class Wav2VecMatcher(nn.Module): ) self.decoder_start_embedding = nn.Parameter(torch.randn(1,1,model_dim)) self.decoder_stop_embedding = nn.Parameter(torch.randn(1,model_dim)) - self.w2v_query_encoder = nn.Linear(WAV2VEC_CHANNELS, model_dim) - self.w2v_value_encoder = nn.Linear(WAV2VEC_CHANNELS, model_dim) + self.w2v_query_encoder = ml.Linear(WAV2VEC_CHANNELS, model_dim) + self.w2v_value_encoder = ml.Linear(WAV2VEC_CHANNELS, model_dim) self.decoder = CheckpointedXTransformer( max_seq_len=-1, # Should be unused use_pos_emb=False, diff --git a/codes/models/classifiers/cifar_resnet.py b/codes/models/classifiers/cifar_resnet.py index ceb78064..86a7af0e 100644 --- a/codes/models/classifiers/cifar_resnet.py +++ b/codes/models/classifiers/cifar_resnet.py @@ -10,6 +10,7 @@ import torch import torch.nn as nn +import torch_intermediary as ml from trainer.networks import register_model @@ -98,7 +99,7 @@ class ResNet(nn.Module): self.conv4_x = self._make_layer(block, 128, num_block[2], 2) self.conv5_x = self._make_layer(block, 256, num_block[3], 2) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) - self.fc = nn.Linear(256 * block.expansion, num_classes) + self.fc = ml.Linear(256 * block.expansion, num_classes) def _make_layer(self, block, out_channels, num_blocks, stride): """make resnet layers(by layer i didnt mean this 'layer' was the diff --git a/codes/models/classifiers/resnet_with_checkpointing.py b/codes/models/classifiers/resnet_with_checkpointing.py index 552d1328..f6f6c5e7 100644 --- a/codes/models/classifiers/resnet_with_checkpointing.py +++ b/codes/models/classifiers/resnet_with_checkpointing.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn from torchvision.models.resnet import BasicBlock, Bottleneck import torchvision +import torch_intermediary as ml __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', @@ -194,5 +195,5 @@ def wide_resnet101_2(pretrained=False, progress=True, **kwargs): def register_resnet50(opt_net, opt): model = resnet50(pretrained=opt_net['pretrained']) if opt_net['custom_head_logits']: - model.fc = nn.Linear(512 * 4, opt_net['custom_head_logits']) + model.fc = ml.Linear(512 * 4, opt_net['custom_head_logits']) return model diff --git a/codes/models/classifiers/twin_cifar_resnet.py b/codes/models/classifiers/twin_cifar_resnet.py index 6aa1f938..bf4d1214 100644 --- a/codes/models/classifiers/twin_cifar_resnet.py +++ b/codes/models/classifiers/twin_cifar_resnet.py @@ -11,6 +11,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +import torch_intermediary as ml from trainer.networks import register_model @@ -101,7 +102,7 @@ class ResNet(nn.Module): self.conv4_x = self._make_layer(block, 128, num_block[2], 2) self.conv5_x = self._make_layer(block, 256, num_block[3], 2) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) - self.fc = nn.Linear(256 * block.expansion, num_classes) + self.fc = ml.Linear(256 * block.expansion, num_classes) def _make_layer(self, block, out_channels, num_blocks, stride): """make resnet layers(by layer i didnt mean this 'layer' was the diff --git a/codes/models/classifiers/weighted_conv_resnet.py b/codes/models/classifiers/weighted_conv_resnet.py index f69c2cb4..5b5beaf6 100644 --- a/codes/models/classifiers/weighted_conv_resnet.py +++ b/codes/models/classifiers/weighted_conv_resnet.py @@ -11,6 +11,7 @@ __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', from models.vqvae.scaled_weight_conv import ScaledWeightConv from trainer.networks import register_model from utils.util import checkpoint +import torch_intermediary as ml model_urls = { 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', @@ -213,7 +214,7 @@ class ResNet(nn.Module): self.layer4 = self._make_layer(block, 512, layers[3], breadth, stride=2, dilate=replace_stride_with_dilation[2]) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) - self.fc = nn.Linear(512 * block.expansion, num_classes) + self.fc = ml.Linear(512 * block.expansion, num_classes) for m in self.modules(): if isinstance(m, ScaledWeightConv): diff --git a/codes/models/classifiers/wide_kernel_vgg.py b/codes/models/classifiers/wide_kernel_vgg.py index d4af0ee0..ab26fb0f 100644 --- a/codes/models/classifiers/wide_kernel_vgg.py +++ b/codes/models/classifiers/wide_kernel_vgg.py @@ -3,7 +3,7 @@ import torch.nn as nn from trainer.networks import register_model from utils.util import opt_get - +import torch_intermediary as ml class WideKernelVgg(nn.Module): def __init__(self, nf=64, num_classes=2): @@ -49,9 +49,9 @@ class WideKernelVgg(nn.Module): nn.ReLU(), nn.MaxPool2d(kernel_size=2), nn.Flatten(), - nn.Linear(nf * 8 * 4 * 2, 100), + ml.Linear(nf * 8 * 4 * 2, 100), nn.ReLU(), - nn.Linear(100, num_classes) + ml.Linear(100, num_classes) ) # These normalization constants should be derived experimentally. diff --git a/codes/models/clip/clvp.py b/codes/models/clip/clvp.py index 7f4a3461..9bfc7655 100644 --- a/codes/models/clip/clvp.py +++ b/codes/models/clip/clvp.py @@ -10,6 +10,7 @@ from models.arch_util import AttentionBlock from models.lucidrains.x_transformers import ContinuousTransformerWrapper, Encoder from trainer.networks import register_model from utils.util import opt_get, checkpoint +import torch_intermediary as ml def exists(val): @@ -58,7 +59,8 @@ class CollapsingTransformer(nn.Module): class ConvFormatEmbedding(nn.Module): def __init__(self, *args, **kwargs): super().__init__() - self.emb = nn.Embedding(*args, **kwargs) + # nn.Embedding + self.emb = ml.Embedding(*args, **kwargs) def forward(self, x): y = self.emb(x) @@ -98,9 +100,10 @@ class CLVP(nn.Module): self.masked_conditioning_latent = nn.Parameter(torch.randn(1,model_dim*2), requires_grad=True) self.mask_conditioning_percentage = mask_conditioning_percentage - self.text_emb = nn.Embedding(num_text_tokens, model_dim) + # nn.Embedding + self.text_emb = ml.Embedding(num_text_tokens, model_dim) self.text_transformer = CollapsingTransformer(model_dim, latent_dim, transformer_heads, dropout, text_enc_depth, text_mask_percentage, use_rms_scaleshift_norm=True) - self.to_text_latent = nn.Linear(latent_dim, latent_dim, bias=False) + self.to_text_latent = ml.Linear(latent_dim, latent_dim, bias=False) self.distributed_collect = distributed_collect if mel_codes is None: @@ -108,7 +111,7 @@ class CLVP(nn.Module): else: self.speech_emb = ConvFormatEmbedding(mel_codes, model_dim) self.speech_transformer = CollapsingTransformer(model_dim, latent_dim, transformer_heads, dropout, speech_enc_depth, speech_mask_percentage) - self.to_speech_latent = nn.Linear(latent_dim, latent_dim, bias=False) + self.to_speech_latent = ml.Linear(latent_dim, latent_dim, bias=False) def get_grad_norm_parameter_groups(self): return { diff --git a/codes/models/clip/contrastive_audio.py b/codes/models/clip/contrastive_audio.py index 6bd77e55..85edfec6 100644 --- a/codes/models/clip/contrastive_audio.py +++ b/codes/models/clip/contrastive_audio.py @@ -9,6 +9,7 @@ from models.arch_util import AttentionBlock from models.lucidrains.x_transformers import ContinuousTransformerWrapper, Encoder from trainer.networks import register_model from utils.util import opt_get, checkpoint +import torch_intermediary as ml def exists(val): @@ -178,7 +179,8 @@ class CollapsingTransformer(nn.Module): class ConvFormatEmbedding(nn.Module): def __init__(self, *args, **kwargs): super().__init__() - self.emb = nn.Embedding(*args, **kwargs) + # nn.Embedding + self.emb = ml.Embedding(*args, **kwargs) def forward(self, x): y = self.emb(x) @@ -203,8 +205,8 @@ class ContrastiveAudio(nn.Module): self.emb = nn.Sequential(nn.Conv1d(mel_channels, model_dim // 2, kernel_size=5, stride=2, padding=2), nn.Conv1d(model_dim//2, model_dim, kernel_size=3, stride=2, padding=1)) self.transformer = CollapsingTransformer(model_dim, model_dim, transformer_heads, dropout, encoder_depth, mask_percent) - self.to_latent = nn.Linear(latent_dim, latent_dim, bias=False) - self.to_latent2 = nn.Linear(latent_dim, latent_dim, bias=False) + self.to_latent = ml.Linear(latent_dim, latent_dim, bias=False) + self.to_latent2 = ml.Linear(latent_dim, latent_dim, bias=False) self.to_latent2.weight.data = self.to_latent.weight.data self.to_latent2.weight.DO_NOT_TRAIN = True diff --git a/codes/models/clip/cvvp.py b/codes/models/clip/cvvp.py index 2ad7eca6..e8a12b2a 100644 --- a/codes/models/clip/cvvp.py +++ b/codes/models/clip/cvvp.py @@ -10,6 +10,7 @@ from models.arch_util import AttentionBlock from models.lucidrains.x_transformers import ContinuousTransformerWrapper, Encoder from trainer.networks import register_model from utils.util import opt_get, checkpoint +import torch_intermediary as ml def exists(val): @@ -58,7 +59,8 @@ class CollapsingTransformer(nn.Module): class ConvFormatEmbedding(nn.Module): def __init__(self, *args, **kwargs): super().__init__() - self.emb = nn.Embedding(*args, **kwargs) + # nn.Embedding + self.emb = ml.Embedding(*args, **kwargs) def forward(self, x): y = self.emb(x) @@ -86,14 +88,14 @@ class CVVP(nn.Module): self.cond_emb = nn.Sequential(nn.Conv1d(mel_channels, model_dim//2, kernel_size=5, stride=2, padding=2), nn.Conv1d(model_dim//2, model_dim, kernel_size=3, stride=2, padding=1)) self.conditioning_transformer = CollapsingTransformer(model_dim, model_dim, transformer_heads, dropout, conditioning_enc_depth, cond_mask_percentage) - self.to_conditioning_latent = nn.Linear(latent_dim, latent_dim, bias=False) + self.to_conditioning_latent = ml.Linear(latent_dim, latent_dim, bias=False) if mel_codes is None: self.speech_emb = nn.Conv1d(mel_channels, model_dim, kernel_size=5, padding=2) else: self.speech_emb = ConvFormatEmbedding(mel_codes, model_dim) self.speech_transformer = CollapsingTransformer(model_dim, latent_dim, transformer_heads, dropout, speech_enc_depth, speech_mask_percentage) - self.to_speech_latent = nn.Linear(latent_dim, latent_dim, bias=False) + self.to_speech_latent = ml.Linear(latent_dim, latent_dim, bias=False) def get_grad_norm_parameter_groups(self): return { diff --git a/codes/models/clip/mel_text_clip.py b/codes/models/clip/mel_text_clip.py index a8c1dda5..c053547d 100644 --- a/codes/models/clip/mel_text_clip.py +++ b/codes/models/clip/mel_text_clip.py @@ -7,6 +7,7 @@ from torch import einsum from models.lucidrains.dalle.transformer import Transformer from trainer.networks import register_model from utils.util import opt_get +import torch_intermediary as ml def exists(val): @@ -45,17 +46,20 @@ class MelTextCLIP(nn.Module): mel_compression=256, ): super().__init__() - self.text_emb = nn.Embedding(num_text_tokens, dim_text) - self.text_pos_emb = nn.Embedding(text_seq_len, dim_text) + # nn.Embedding + self.text_emb = ml.Embedding(num_text_tokens, dim_text) + # nn.Embedding + self.text_pos_emb = ml.Embedding(text_seq_len, dim_text) self.text_transformer = Transformer(causal=False, seq_len=text_seq_len, dim=dim_text, depth=text_enc_depth, heads=text_heads, rotary_emb=False) - self.to_text_latent = nn.Linear(dim_text, dim_latent, bias=False) + self.to_text_latent = ml.Linear(dim_text, dim_latent, bias=False) self.speech_enc = nn.Conv1d(80, dim_speech, kernel_size=3, padding=1) - self.speech_pos_emb = nn.Embedding(num_speech_tokens, dim_speech) + # nn.Embedding + self.speech_pos_emb = ml.Embedding(num_speech_tokens, dim_speech) self.speech_transformer = Transformer(causal=False, seq_len=speech_seq_len, dim=dim_speech, depth=speech_enc_depth, heads=speech_heads, rotary_emb=False) - self.to_speech_latent = nn.Linear(dim_speech, dim_latent, bias=False) + self.to_speech_latent = ml.Linear(dim_speech, dim_latent, bias=False) self.temperature = nn.Parameter(torch.tensor(1.)) self.text_mask_percentage = text_mask_percentage diff --git a/codes/models/clip/text_cond_clip.py b/codes/models/clip/text_cond_clip.py index 39085202..f221142c 100644 --- a/codes/models/clip/text_cond_clip.py +++ b/codes/models/clip/text_cond_clip.py @@ -7,6 +7,7 @@ from models.audio.tts.unified_voice2 import ConditioningEncoder from models.lucidrains.dalle.transformer import Transformer from trainer.networks import register_model from utils.util import opt_get +import torch_intermediary as ml def exists(val): @@ -45,7 +46,7 @@ class VoiceCondCLIP(nn.Module): self.speech_pos_emb = nn.Embedding(num_speech_tokens, dim_speech) self.speech_transformer = Transformer(causal=False, seq_len=speech_seq_len, dim=dim_speech, depth=speech_enc_depth, heads=speech_heads, rotary_emb=False) - self.to_speech_latent = nn.Linear(dim_speech, dim_latent, bias=False) + self.to_speech_latent = ml.Linear(dim_speech, dim_latent, bias=False) self.temperature = nn.Parameter(torch.tensor(1.)) self.voice_mask_percentage = voice_mask_percentage diff --git a/codes/models/clip/text_voice_clip.py b/codes/models/clip/text_voice_clip.py index 297e71dc..26bac0c9 100644 --- a/codes/models/clip/text_voice_clip.py +++ b/codes/models/clip/text_voice_clip.py @@ -11,6 +11,7 @@ from models.audio.tts.unet_diffusion_tts7 import CheckpointedXTransformerEncoder from models.lucidrains.dalle.transformer import Transformer from trainer.networks import register_model from utils.util import opt_get +import torch_intermediary as ml def exists(val): @@ -53,11 +54,13 @@ class VoiceCLIP(nn.Module): distributed_collect=False, ): super().__init__() - self.text_emb = nn.Embedding(num_text_tokens, dim_text) - self.to_text_latent = nn.Linear(dim_text, dim_latent, bias=False) + # nn.Embedding + self.text_emb = ml.Embedding(num_text_tokens, dim_text) + self.to_text_latent = ml.Linear(dim_text, dim_latent, bias=False) - self.speech_emb = nn.Embedding(num_speech_tokens, dim_speech) - self.to_speech_latent = nn.Linear(dim_speech, dim_latent, bias=False) + # nn.Embedding + self.speech_emb = ml.Embedding(num_speech_tokens, dim_speech) + self.to_speech_latent = ml.Linear(dim_speech, dim_latent, bias=False) if use_xformers: self.text_transformer = CheckpointedXTransformerEncoder( @@ -105,8 +108,10 @@ class VoiceCLIP(nn.Module): self.min_mel_size = min_mel_size self.distributed_collect = distributed_collect if not use_xformers: - self.text_pos_emb = nn.Embedding(text_seq_len, dim_text) - self.speech_pos_emb = nn.Embedding(num_speech_tokens, dim_speech) + # nn.Embedding + self.text_pos_emb = ml.Embedding(text_seq_len, dim_text) + # nn.Embedding + self.speech_pos_emb = ml.Embedding(num_speech_tokens, dim_speech) def embed_text(self, text): text_mask = torch.ones_like(text.float()).bool() diff --git a/codes/models/diffusion/nn.py b/codes/models/diffusion/nn.py index 50203d75..201fadb9 100644 --- a/codes/models/diffusion/nn.py +++ b/codes/models/diffusion/nn.py @@ -6,6 +6,7 @@ import math import torch as th import torch.nn as nn +import torch_intermediary as ml # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. @@ -36,7 +37,7 @@ def linear(*args, **kwargs): """ Create a linear module. """ - return nn.Linear(*args, **kwargs) + return ml.Linear(*args, **kwargs) def avg_pool_nd(dims, *args, **kwargs): diff --git a/codes/models/diffusion/rrdb_diffusion.py b/codes/models/diffusion/rrdb_diffusion.py index 9deab2f3..933fe467 100644 --- a/codes/models/diffusion/rrdb_diffusion.py +++ b/codes/models/diffusion/rrdb_diffusion.py @@ -6,6 +6,7 @@ from models.arch_util import ConvGnLelu, default_init_weights, make_layer from models.diffusion.nn import timestep_embedding from trainer.networks import register_model from utils.util import checkpoint +import torch_intermediary as ml # Conditionally uses torch's checkpoint functionality if it is enabled in the opt file. @@ -28,7 +29,7 @@ class ResidualDenseBlock(nn.Module): self.first_conv = ConvGnLelu(mid_channels, mid_channels, activation=True, norm=False, bias=True) self.emb_layers = nn.Sequential( nn.SiLU(), - nn.Linear( + ml.Linear( mid_channels*4, mid_channels, ), @@ -143,9 +144,9 @@ class RRDBNet(nn.Module): # Guided diffusion uses a time embedding. time_embed_dim = mid_channels * 4 self.time_embed = nn.Sequential( - nn.Linear(mid_channels, time_embed_dim), + ml.Linear(mid_channels, time_embed_dim), nn.SiLU(), - nn.Linear(time_embed_dim, time_embed_dim), + ml.Linear(time_embed_dim, time_embed_dim), ) self.body = make_layer( diff --git a/codes/models/diffusion/unet_diffusion.py b/codes/models/diffusion/unet_diffusion.py index e2ed97b7..6ebbae02 100644 --- a/codes/models/diffusion/unet_diffusion.py +++ b/codes/models/diffusion/unet_diffusion.py @@ -20,6 +20,7 @@ from models.diffusion.nn import ( ) from trainer.networks import register_model from utils.util import checkpoint +import torch_intermediary as ml class AttentionPool2d(nn.Module): @@ -515,7 +516,8 @@ class UNetModel(nn.Module): ) if self.num_classes is not None: - self.label_emb = nn.Embedding(num_classes, time_embed_dim) + # nn.Embedding + self.label_emb = ml.Embedding(num_classes, time_embed_dim) self.use_raw_y_as_embedding = use_raw_y_as_embedding assert not ((self.num_classes is not None) and use_raw_y_as_embedding) # These are mutually-exclusive. @@ -867,16 +869,16 @@ class EncoderUNetModel(nn.Module): ) elif pool == "spatial": self.out = nn.Sequential( - nn.Linear(self._feature_size, 2048), + ml.Linear(self._feature_size, 2048), nn.ReLU(), - nn.Linear(2048, self.out_channels), + ml.Linear(2048, self.out_channels), ) elif pool == "spatial_v2": self.out = nn.Sequential( - nn.Linear(self._feature_size, 2048), + ml.Linear(self._feature_size, 2048), normalization(2048), nn.SiLU(), - nn.Linear(2048, self.out_channels), + ml.Linear(2048, self.out_channels), ) else: raise NotImplementedError(f"Unexpected {pool} pooling") diff --git a/codes/models/diffusion/unet_latent_guide.py b/codes/models/diffusion/unet_latent_guide.py index 41dd85a8..e298a900 100644 --- a/codes/models/diffusion/unet_latent_guide.py +++ b/codes/models/diffusion/unet_latent_guide.py @@ -26,6 +26,7 @@ from models.diffusion.nn import ( ) from trainer.networks import register_model from utils.util import checkpoint +import torch_intermediary as ml class AttentionPool2d(nn.Module): @@ -476,7 +477,8 @@ class UNetModel(nn.Module): ) if self.num_classes is not None: - self.label_emb = nn.Embedding(num_classes, time_embed_dim) + # nn.Embedding + self.label_emb = ml.Embedding(num_classes, time_embed_dim) self.input_blocks = nn.ModuleList( [ @@ -736,7 +738,7 @@ class ResNetEncoder(nn.Module): dilate=replace_stride_with_dilation[2]) f=512 self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) - self.fc = nn.Linear(f * block.expansion, output_dim) + self.fc = ml.Linear(f * block.expansion, output_dim) for m in self.modules(): if isinstance(m, nn.Conv2d): diff --git a/codes/models/image_generation/discriminator_vgg_arch.py b/codes/models/image_generation/discriminator_vgg_arch.py index 234272de..af44eca2 100644 --- a/codes/models/image_generation/discriminator_vgg_arch.py +++ b/codes/models/image_generation/discriminator_vgg_arch.py @@ -5,6 +5,7 @@ import torch.nn.functional as F from trainer.networks import register_model from utils.util import checkpoint, opt_get +import torch_intermediary as ml class Discriminator_VGG_128(nn.Module): @@ -46,8 +47,8 @@ class Discriminator_VGG_128(nn.Module): input_img_factor = input_img_factor // 2 final_nf = nf * 16 - self.linear1 = nn.Linear(final_nf * 4 * input_img_factor * 4 * input_img_factor, 100) - self.linear2 = nn.Linear(100, 1) + self.linear1 = ml.Linear(final_nf * 4 * input_img_factor * 4 * input_img_factor, 100) + self.linear2 = ml.Linear(100, 1) # activation function self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) @@ -129,8 +130,8 @@ class Discriminator_VGG_128_GN(nn.Module): # activation function self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) - self.linear1 = nn.Linear(int(final_nf * 4 * input_img_factor * 4 * input_img_factor), 100) - self.linear2 = nn.Linear(100, 1) + self.linear1 = ml.Linear(int(final_nf * 4 * input_img_factor * 4 * input_img_factor), 100) + self.linear2 = ml.Linear(100, 1) def compute_body(self, x): fea = self.lrelu(self.conv0_0(x)) @@ -219,8 +220,8 @@ class DiscriminatorVGG448GN(nn.Module): self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) final_nf = nf * 8 - self.linear1 = nn.Linear(int(final_nf * 7 * 7), 100) - self.linear2 = nn.Linear(100, 1) + self.linear1 = ml.Linear(int(final_nf * 7 * 7), 100) + self.linear2 = ml.Linear(100, 1) # Assign all new heads to the new param group.2 for m in [self.convn1_0, self.convn1_1, self.bnn1_1, self.conv0_0_new, self.bn0_0]: diff --git a/codes/models/image_generation/srflow/module_util.py b/codes/models/image_generation/srflow/module_util.py index ca5d7fa9..f50198dd 100644 --- a/codes/models/image_generation/srflow/module_util.py +++ b/codes/models/image_generation/srflow/module_util.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn import torch.nn.init as init import torch.nn.functional as F +import torch_intermediary as ml def initialize_weights(net_l, scale=1): @@ -14,7 +15,7 @@ def initialize_weights(net_l, scale=1): m.weight.data *= scale # for residual block if m.bias is not None: m.bias.data.zero_() - elif isinstance(m, nn.Linear): + elif isinstance(m, ml.Linear): init.kaiming_normal_(m.weight, a=0, mode='fan_in') m.weight.data *= scale if m.bias is not None: diff --git a/codes/models/image_generation/stylegan/stylegan2_lucidrains.py b/codes/models/image_generation/stylegan/stylegan2_lucidrains.py index 0f65a7ef..6016fb9e 100644 --- a/codes/models/image_generation/stylegan/stylegan2_lucidrains.py +++ b/codes/models/image_generation/stylegan/stylegan2_lucidrains.py @@ -28,6 +28,7 @@ except: APEX_AVAILABLE = False assert torch.cuda.is_available(), 'You need to have an Nvidia GPU with CUDA installed.' +import torch_intermediary as ml num_cores = multiprocessing.cpu_count() @@ -351,7 +352,7 @@ class RGBBlock(nn.Module): def __init__(self, latent_dim, input_channel, upsample, rgba=False): super().__init__() self.input_channel = input_channel - self.to_style = nn.Linear(latent_dim, input_channel) + self.to_style = ml.Linear(latent_dim, input_channel) out_filters = 3 if not rgba else 4 self.conv = Conv2DMod(input_channel, out_filters, 1, demod=False) @@ -489,16 +490,16 @@ class GeneratorBlockWithStructure(nn.Module): # Uses stylegan1 style blocks for injecting structural latent. self.conv0 = EqualConv2d(input_channels, filters, 3, padding=1) - self.to_noise0 = nn.Linear(1, filters) + self.to_noise0 = ml.Linear(1, filters) self.noise0 = equal_lr(NoiseInjection(filters)) self.adain0 = AdaptiveInstanceNorm(filters, latent_dim) - self.to_style1 = nn.Linear(latent_dim, filters) - self.to_noise1 = nn.Linear(1, filters) + self.to_style1 = ml.Linear(latent_dim, filters) + self.to_noise1 = ml.Linear(1, filters) self.conv1 = Conv2DMod(filters, filters, 3) - self.to_style2 = nn.Linear(latent_dim, filters) - self.to_noise2 = nn.Linear(1, filters) + self.to_style2 = ml.Linear(latent_dim, filters) + self.to_noise2 = ml.Linear(1, filters) self.conv2 = Conv2DMod(filters, filters, 3) self.activation = leaky_relu() @@ -540,12 +541,12 @@ class GeneratorBlock(nn.Module): self.structure_conv = nn.Conv2d(3, input_channels, 3, padding=1) input_channels = input_channels * 2 - self.to_style1 = nn.Linear(latent_dim, input_channels) - self.to_noise1 = nn.Linear(1, filters) + self.to_style1 = ml.Linear(latent_dim, input_channels) + self.to_noise1 = ml.Linear(1, filters) self.conv1 = Conv2DMod(input_channels, filters, 3) - self.to_style2 = nn.Linear(latent_dim, filters) - self.to_noise2 = nn.Linear(1, filters) + self.to_style2 = ml.Linear(latent_dim, filters) + self.to_noise2 = ml.Linear(1, filters) self.conv2 = Conv2DMod(filters, filters, 3) self.activation = leaky_relu() @@ -724,7 +725,7 @@ class StyleGan2GeneratorWithLatent(nn.Module): def _init_weights(self): for m in self.modules(): - if type(m) in {nn.Conv2d, nn.Linear} and hasattr(m, 'weight'): + if type(m) in {nn.Conv2d, ml.Linear} and hasattr(m, 'weight'): nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in', nonlinearity='leaky_relu') for block in self.gen.blocks: @@ -804,7 +805,7 @@ class StyleGan2Discriminator(nn.Module): self.final_conv = nn.Conv2d(chan_last, chan_last, 3, padding=1) self.flatten = Flatten() - self.to_logit = nn.Linear(latent_dim, 1) + self.to_logit = ml.Linear(latent_dim, 1) self._init_weights() @@ -836,7 +837,7 @@ class StyleGan2Discriminator(nn.Module): def _init_weights(self): for m in self.modules(): - if type(m) in {nn.Conv2d, nn.Linear}: + if type(m) in {nn.Conv2d, ml.Linear}: nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in', nonlinearity='leaky_relu') diff --git a/codes/models/image_latents/byol/byol_model_wrapper.py b/codes/models/image_latents/byol/byol_model_wrapper.py index cf3b6ea0..6c8bb3e3 100644 --- a/codes/models/image_latents/byol/byol_model_wrapper.py +++ b/codes/models/image_latents/byol/byol_model_wrapper.py @@ -12,6 +12,7 @@ from torch import nn from data.images.byol_attachment import RandomApply from trainer.networks import register_model, create_model from utils.util import checkpoint, opt_get +import torch_intermediary as ml def default(val, def_val): @@ -78,10 +79,10 @@ class MLP(nn.Module): def __init__(self, dim, projection_size, hidden_size=4096): super().__init__() self.net = nn.Sequential( - nn.Linear(dim, hidden_size), + ml.Linear(dim, hidden_size), nn.BatchNorm1d(hidden_size), nn.ReLU(inplace=True), - nn.Linear(hidden_size, projection_size) + ml.Linear(hidden_size, projection_size) ) def forward(self, x): @@ -103,10 +104,10 @@ class StructuralMLP(nn.Module): nn.BatchNorm2d(c), nn.ReLU(inplace=True), nn.Flatten(), - nn.Linear(flattened_dim, hidden_size), + ml.Linear(flattened_dim, hidden_size), nn.BatchNorm1d(hidden_size), nn.ReLU(inplace=True), - nn.Linear(hidden_size, projection_size) + ml.Linear(hidden_size, projection_size) ) def forward(self, x): diff --git a/codes/models/image_latents/fixup_resnet/DiscriminatorResnet_arch.py b/codes/models/image_latents/fixup_resnet/DiscriminatorResnet_arch.py index 6991b0fc..2c2a8fcd 100644 --- a/codes/models/image_latents/fixup_resnet/DiscriminatorResnet_arch.py +++ b/codes/models/image_latents/fixup_resnet/DiscriminatorResnet_arch.py @@ -1,6 +1,7 @@ import torch import torch.nn as nn import numpy as np +import torch_intermediary as ml __all__ = ['FixupResNet', 'fixup_resnet18', 'fixup_resnet34', 'fixup_resnet50', 'fixup_resnet101', 'fixup_resnet152'] @@ -108,8 +109,8 @@ class FixupResNet(nn.Module): self.layer4 = self._make_layer(block, num_filters*8, layers[3], stride=2) self.bias2 = nn.Parameter(torch.zeros(1)) reduced_img_sz = int(input_img_size / 32) - self.fc1 = nn.Linear(num_filters * 8 * reduced_img_sz * reduced_img_sz, 100) - self.fc2 = nn.Linear(100, num_classes) + self.fc1 = ml.Linear(num_filters * 8 * reduced_img_sz * reduced_img_sz, 100) + self.fc2 = ml.Linear(100, num_classes) for m in self.modules(): if isinstance(m, FixupBasicBlock): @@ -124,7 +125,7 @@ class FixupResNet(nn.Module): if m.downsample is not None: nn.init.normal_(m.downsample.weight, mean=0, std=np.sqrt(2 / (m.downsample.weight.shape[0] * np.prod(m.downsample.weight.shape[2:])))) ''' - elif isinstance(m, nn.Linear): + elif isinstance(m, ml.Linear): nn.init.constant_(m.weight, 0) nn.init.constant_(m.bias, 0)''' diff --git a/codes/models/image_latents/vit_latent.py b/codes/models/image_latents/vit_latent.py index 0243e8a7..b38671be 100644 --- a/codes/models/image_latents/vit_latent.py +++ b/codes/models/image_latents/vit_latent.py @@ -5,6 +5,7 @@ import torch.nn.functional as F from models.arch_util import ResBlock from models.lucidrains.x_transformers import Encoder from trainer.networks import register_model +import torch_intermediary as ml class VitLatent(nn.Module): @@ -31,10 +32,10 @@ class VitLatent(nn.Module): do_checkpointing=True ) - self.mlp = nn.Sequential(nn.Linear(hidden_dim, hidden_dim*2), + self.mlp = nn.Sequential(ml.Linear(hidden_dim, hidden_dim*2), nn.BatchNorm1d(hidden_dim*2), nn.ReLU(inplace=True), - nn.Linear(hidden_dim*2, hidden_dim)) + ml.Linear(hidden_dim*2, hidden_dim)) def provide_ema(self, ema): self.ema = ema diff --git a/codes/models/lucidrains/dalle/attention.py b/codes/models/lucidrains/dalle/attention.py index c3b52cc4..09ec3978 100644 --- a/codes/models/lucidrains/dalle/attention.py +++ b/codes/models/lucidrains/dalle/attention.py @@ -7,6 +7,7 @@ import torch.nn.functional as F from einops import rearrange, repeat from rotary_embedding_torch import apply_rotary_emb +import torch_intermediary as ml # helpers @@ -47,9 +48,9 @@ class Attention(nn.Module): self.stable = stable self.causal = causal - self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False) + self.to_qkv = ml.Linear(dim, inner_dim * 3, bias = False) self.to_out = nn.Sequential( - nn.Linear(inner_dim, dim), + ml.Linear(inner_dim, dim), nn.Dropout(dropout) ) @@ -102,10 +103,10 @@ class SparseConvCausalAttention(nn.Module): self.stable = stable - self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False) + self.to_qkv = ml.Linear(dim, inner_dim * 3, bias = False) self.to_out = nn.Sequential( - nn.Linear(inner_dim, dim), + ml.Linear(inner_dim, dim), nn.Dropout(dropout) ) @@ -222,10 +223,10 @@ class SparseAxialCausalAttention(nn.Module): self.stable = stable - self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False) + self.to_qkv = ml.Linear(dim, inner_dim * 3, bias = False) self.to_out = nn.Sequential( - nn.Linear(inner_dim, dim), + ml.Linear(inner_dim, dim), nn.Dropout(dropout) ) diff --git a/codes/models/lucidrains/dalle/transformer.py b/codes/models/lucidrains/dalle/transformer.py index 389e6849..357ce1ad 100644 --- a/codes/models/lucidrains/dalle/transformer.py +++ b/codes/models/lucidrains/dalle/transformer.py @@ -11,6 +11,7 @@ from models.lucidrains.dalle.attention import Attention, SparseAttention, Sparse from rotary_embedding_torch import RotaryEmbedding, broadcat from g_mlp_pytorch import gMLPBlock +import torch_intermediary as ml # helpers @@ -78,10 +79,10 @@ class FeedForward(nn.Module): def __init__(self, dim, dropout = 0., mult = 4.): super().__init__() self.net = nn.Sequential( - nn.Linear(dim, dim * mult * 2), + ml.Linear(dim, dim * mult * 2), GEGLU(), nn.Dropout(dropout), - nn.Linear(dim * mult, dim) + ml.Linear(dim * mult, dim) ) def forward(self, x): diff --git a/codes/models/lucidrains/performer/performer_pytorch.py b/codes/models/lucidrains/performer/performer_pytorch.py index 2d618f11..98ce769f 100644 --- a/codes/models/lucidrains/performer/performer_pytorch.py +++ b/codes/models/lucidrains/performer/performer_pytorch.py @@ -21,6 +21,7 @@ try: APEX_AVAILABLE = True except: APEX_AVAILABLE = False +import torch_intermediary as ml # helpers @@ -356,10 +357,10 @@ class FeedForward(nn.Module): activation = default(activation, nn.GELU) self.glu = glu - self.w1 = nn.Linear(dim, dim * mult * (2 if glu else 1)) + self.w1 = ml.Linear(dim, dim * mult * (2 if glu else 1)) self.act = activation() self.dropout = nn.Dropout(dropout) - self.w2 = nn.Linear(dim * mult, dim) + self.w2 = ml.Linear(dim * mult, dim) def forward(self, x, **kwargs): if not self.glu: @@ -401,10 +402,10 @@ class Attention(nn.Module): self.global_heads = heads - local_heads self.local_attn = LocalAttention(window_size = local_window_size, causal = causal, autopad = True, dropout = dropout, look_forward = int(not causal), rel_pos_emb_config = (dim_head, local_heads)) if local_heads > 0 else None - self.to_q = nn.Linear(dim, inner_dim, bias = qkv_bias) - self.to_k = nn.Linear(dim, inner_dim, bias = qkv_bias) - self.to_v = nn.Linear(dim, inner_dim, bias = qkv_bias) - self.to_out = nn.Linear(inner_dim, dim, bias = attn_out_bias) + self.to_q = ml.Linear(dim, inner_dim, bias = qkv_bias) + self.to_k = ml.Linear(dim, inner_dim, bias = qkv_bias) + self.to_v = ml.Linear(dim, inner_dim, bias = qkv_bias) + self.to_out = ml.Linear(inner_dim, dim, bias = attn_out_bias) self.dropout = nn.Dropout(dropout) def forward(self, x, pos_emb = None, context = None, mask = None, context_mask = None, **kwargs): @@ -458,7 +459,8 @@ class CrossAttention(Attention): class AbsolutePositionalEmbedding(nn.Module): def __init__(self, dim, max_seq_len): super().__init__() - self.emb = nn.Embedding(max_seq_len, dim) + # nn.Embedding + self.emb = ml.Embedding(max_seq_len, dim) def forward(self, x): t = torch.arange(x.shape[1], device=x.device) @@ -619,7 +621,8 @@ class PerformerLM(nn.Module): local_attn_heads = cast_tuple(local_attn_heads) self.max_seq_len = max_seq_len - self.token_emb = nn.Embedding(num_tokens, dim) + # nn.Embedding + self.token_emb = ml.Embedding(num_tokens, dim) if rotary_position_emb: self.pos_emb = FixedPositionalEmbedding(dim, max_seq_len) @@ -636,7 +639,7 @@ class PerformerLM(nn.Module): self.performer = Performer(dim, depth, heads, dim_head, local_attn_heads, local_window_size, causal, ff_mult, nb_features, feature_redraw_interval, reversible, ff_chunks, generalized_attention, kernel_fn, use_scalenorm, use_rezero, ff_glu, ff_dropout, attn_dropout, cross_attend, no_projection, auto_check_redraw, qkv_bias, attn_out_bias, shift_tokens) self.norm = nn.LayerNorm(dim) - self.to_out = nn.Linear(dim, num_tokens) if not tie_embed else None + self.to_out = ml.Linear(dim, num_tokens) if not tie_embed else None def check_redraw_projections(self): self.performer.check_redraw_projections() diff --git a/codes/models/lucidrains/vq.py b/codes/models/lucidrains/vq.py index 4b1019ec..13a4b2ae 100644 --- a/codes/models/lucidrains/vq.py +++ b/codes/models/lucidrains/vq.py @@ -8,6 +8,7 @@ from torch.cuda.amp import autocast from einops import rearrange, repeat from contextlib import contextmanager +import torch_intermediary as ml def par(t, nm): @@ -355,9 +356,9 @@ class VectorQuantize(nn.Module): codebook_dim = default(codebook_dim, dim) requires_projection = codebook_dim != dim - self.project_in = nn.Linear(dim, codebook_dim) if requires_projection \ + self.project_in = ml.Linear(dim, codebook_dim) if requires_projection \ else nn.Identity() - self.project_out = nn.Linear(codebook_dim, dim) if requires_projection \ + self.project_out = ml.Linear(codebook_dim, dim) if requires_projection \ else nn.Identity() self.eps = eps diff --git a/codes/models/lucidrains/x_transformers.py b/codes/models/lucidrains/x_transformers.py index 32ad19e7..a49af93d 100644 --- a/codes/models/lucidrains/x_transformers.py +++ b/codes/models/lucidrains/x_transformers.py @@ -11,6 +11,7 @@ from einops import rearrange, repeat, reduce from einops.layers.torch import Rearrange from torch.utils.checkpoint import checkpoint +import torch_intermediary as ml DEFAULT_DIM_HEAD = 64 @@ -125,7 +126,8 @@ class AbsolutePositionalEmbedding(nn.Module): def __init__(self, dim, max_seq_len): super().__init__() self.scale = dim ** -0.5 - self.emb = nn.Embedding(max_seq_len, dim) + # nn.Embedding + self.emb = ml.Embedding(max_seq_len, dim) def forward(self, x): n = torch.arange(x.shape[1], device=x.device) @@ -154,7 +156,8 @@ class RelativePositionBias(nn.Module): self.causal = causal self.num_buckets = num_buckets self.max_distance = max_distance - self.relative_attention_bias = nn.Embedding(num_buckets, heads) + # nn.Embedding + self.relative_attention_bias = ml.Embedding(num_buckets, heads) @staticmethod def _relative_position_bucket(relative_position, causal=True, num_buckets=32, max_distance=128): @@ -360,7 +363,7 @@ class RMSScaleShiftNorm(nn.Module): self.cdim = 1 self.pdim = -1 else: - self.scale_shift_process = nn.Linear(embed_dim, dim * 2, bias=bias) + self.scale_shift_process = ml.Linear(embed_dim, dim * 2, bias=bias) self.cdim = -1 self.pdim = 1 @@ -447,7 +450,7 @@ class GLU(nn.Module): def __init__(self, dim_in, dim_out, activation): super().__init__() self.act = activation - self.proj = nn.Linear(dim_in, dim_out * 2) + self.proj = ml.Linear(dim_in, dim_out * 2) def forward(self, x): x, gate = self.proj(x).chunk(2, dim=-1) @@ -472,7 +475,7 @@ class FeedForward(nn.Module): activation = ReluSquared() if relu_squared else nn.GELU() project_in = nn.Sequential( - nn.Linear(dim, inner_dim), + ml.Linear(dim, inner_dim), activation ) if not glu else GLU(dim, inner_dim, activation) @@ -480,7 +483,7 @@ class FeedForward(nn.Module): project_in, nn.LayerNorm(inner_dim) if post_act_ln else nn.Identity(), nn.Dropout(dropout), - nn.Linear(inner_dim, dim_out) + ml.Linear(inner_dim, dim_out) ) # init last linear layer to 0 @@ -535,16 +538,16 @@ class Attention(nn.Module): qk_dim = int(collab_compression * qk_dim) self.collab_mixing = nn.Parameter(torch.randn(heads, qk_dim)) - self.to_q = nn.Linear(dim, qk_dim, bias=False) - self.to_k = nn.Linear(dim, qk_dim, bias=False) - self.to_v = nn.Linear(dim, v_dim, bias=False) + self.to_q = ml.Linear(dim, qk_dim, bias=False) + self.to_k = ml.Linear(dim, qk_dim, bias=False) + self.to_v = ml.Linear(dim, v_dim, bias=False) self.dropout = nn.Dropout(dropout) # add GLU gating for aggregated values, from alphafold2 self.to_v_gate = None if gate_values: - self.to_v_gate = nn.Linear(dim, v_dim) + self.to_v_gate = ml.Linear(dim, v_dim) nn.init.constant_(self.to_v_gate.weight, 0) nn.init.constant_(self.to_v_gate.bias, 1) @@ -581,7 +584,7 @@ class Attention(nn.Module): # attention on attention self.attn_on_attn = on_attn out_dim = default(out_dim, dim) - self.to_out = nn.Sequential(nn.Linear(v_dim, out_dim * 2), nn.GLU()) if on_attn else nn.Linear(v_dim, out_dim) + self.to_out = nn.Sequential(ml.Linear(v_dim, out_dim * 2), nn.GLU()) if on_attn else ml.Linear(v_dim, out_dim) self.rel_pos_bias = rel_pos_bias if rel_pos_bias: @@ -1077,7 +1080,7 @@ class ViTransformerWrapper(nn.Module): self.patch_size = patch_size self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim)) - self.patch_to_embedding = nn.Linear(patch_dim, dim) + self.patch_to_embedding = ml.Linear(patch_dim, dim) self.cls_token = nn.Parameter(torch.randn(1, 1, dim)) self.dropout = nn.Dropout(emb_dropout) @@ -1135,18 +1138,19 @@ class TransformerWrapper(nn.Module): self.max_mem_len = max_mem_len self.shift_mem_down = shift_mem_down - self.token_emb = nn.Embedding(num_tokens, emb_dim) + # nn.Embedding + self.token_emb = ml.Embedding(num_tokens, emb_dim) self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len) if ( use_pos_emb and not attn_layers.has_pos_emb) else always(0) self.emb_dropout = nn.Dropout(emb_dropout) - self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity() + self.project_emb = ml.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity() self.attn_layers = attn_layers self.norm = nn.LayerNorm(dim) self.init_() - self.to_logits = nn.Linear(dim, num_tokens) if not tie_embedding else lambda t: t @ self.token_emb.weight.t() + self.to_logits = ml.Linear(dim, num_tokens) if not tie_embedding else lambda t: t @ self.token_emb.weight.t() # memory tokens (like [cls]) from Memory Transformers paper num_memory_tokens = default(num_memory_tokens, 0) @@ -1233,12 +1237,12 @@ class ContinuousTransformerWrapper(nn.Module): use_pos_emb and not attn_layers.has_pos_emb) else always(0) self.emb_dropout = nn.Dropout(emb_dropout) - self.project_in = nn.Linear(dim_in, dim) if exists(dim_in) else nn.Identity() + self.project_in = ml.Linear(dim_in, dim) if exists(dim_in) else nn.Identity() self.attn_layers = attn_layers self.norm = nn.LayerNorm(dim) - self.project_out = nn.Linear(dim, dim_out) if exists(dim_out) else nn.Identity() + self.project_out = ml.Linear(dim, dim_out) if exists(dim_out) else nn.Identity() def forward( self, diff --git a/codes/models/vqvae/gumbel_quantizer.py b/codes/models/vqvae/gumbel_quantizer.py index 2d71ec7f..e1d95f26 100644 --- a/codes/models/vqvae/gumbel_quantizer.py +++ b/codes/models/vqvae/gumbel_quantizer.py @@ -4,13 +4,15 @@ import torch.nn.functional as F from torch import einsum from utils.weight_scheduler import LinearDecayWeightScheduler +import torch_intermediary as ml class GumbelQuantizer(nn.Module): def __init__(self, inp_dim, codebook_dim, num_tokens, straight_through=False): super().__init__() self.to_logits = nn.Conv1d(inp_dim, num_tokens, 1) - self.codebook = nn.Embedding(num_tokens, codebook_dim) + # nn.Embedding + self.codebook = ml.Embedding(num_tokens, codebook_dim) self.straight_through = straight_through self.temperature_scheduler = LinearDecayWeightScheduler(10, 5000, .9, 2000) self.step = 0 diff --git a/codes/models/vqvae/vector_quantizer.py b/codes/models/vqvae/vector_quantizer.py index 96015b02..6d057c25 100644 --- a/codes/models/vqvae/vector_quantizer.py +++ b/codes/models/vqvae/vector_quantizer.py @@ -4,6 +4,7 @@ import torch.nn.functional as F from einops import rearrange, repeat from models.arch_util import l2norm, sample_vectors, default, ema_inplace +import torch_intermediary as ml def kmeans(samples, num_clusters, num_iters = 10, use_cosine_sim = False): @@ -184,8 +185,8 @@ class VectorQuantize(nn.Module): codebook_dim = default(codebook_dim, dim) requires_projection = codebook_dim != dim - self.project_in = nn.Linear(dim, codebook_dim) if requires_projection else nn.Identity() - self.project_out = nn.Linear(codebook_dim, dim) if requires_projection else nn.Identity() + self.project_in = ml.Linear(dim, codebook_dim) if requires_projection else nn.Identity() + self.project_out = ml.Linear(codebook_dim, dim) if requires_projection else nn.Identity() self.eps = eps diff --git a/codes/torch_intermediary/__init__.py b/codes/torch_intermediary/__init__.py new file mode 100644 index 00000000..bc8933a6 --- /dev/null +++ b/codes/torch_intermediary/__init__.py @@ -0,0 +1,41 @@ +""" +from bitsandbytes.nn import Linear8bitLt as Linear +from bitsandbytes.nn import StableEmbedding as Embedding +from bitsandbytes.optim.adam import Adam8bit as Adam +from bitsandbytes.optim.adamw import AdamW8bit as AdamW +""" +""" +from torch.nn import Linear +from torch.nn import Embedding +from torch.optim.adam import Adam +from torch.optim.adamw import AdamW +""" + +OVERRIDE_LINEAR = False +OVERRIDE_EMBEDDING = False +OVERRIDE_ADAM = True +OVERRIDE_ADAMW = True +USE_STABLE_EMBEDDING = True + +if OVERRIDE_LINEAR: + from bitsandbytes.nn import Linear8bitLt as Linear +else: + from torch.nn import Linear + +if OVERRIDE_EMBEDDING: + if USE_STABLE_EMBEDDING: + from bitsandbytes.nn import StableEmbedding as Embedding + else: + from bitsandbytes.nn import Embedding as Embedding +else: + from torch.nn import Embedding + +if OVERRIDE_ADAM: + from bitsandbytes.optim.adam import Adam8bit as Adam +else: + from torch.optim.adam import Adam + +if OVERRIDE_ADAMW: + from bitsandbytes.optim.adamw import AdamW8bit as AdamW +else: + from torch.optim.adamw import AdamW \ No newline at end of file diff --git a/codes/trainer/ExtensibleTrainer.py b/codes/trainer/ExtensibleTrainer.py index 3ee81074..0c80ede7 100644 --- a/codes/trainer/ExtensibleTrainer.py +++ b/codes/trainer/ExtensibleTrainer.py @@ -21,6 +21,7 @@ import torchvision.utils as utils from utils.loss_accumulator import LossAccumulator, InfStorageLossAccumulator from utils.util import opt_get, denormalize +import torch_intermediary as ml logger = logging.getLogger('base') @@ -337,7 +338,7 @@ class ExtensibleTrainer(BaseModel): for net in self.networks.values(): for mod in net.modules(): fan_in = -1 - if isinstance(mod, nn.Linear): + if isinstance(mod, ml.Linear): fan_in = mod.weight.data.shape[1] elif isinstance(mod, nn.Conv1d): fan_in = mod.weight.data.shape[0] diff --git a/codes/trainer/feature_model.py b/codes/trainer/feature_model.py index 12e31cb1..c42c9563 100644 --- a/codes/trainer/feature_model.py +++ b/codes/trainer/feature_model.py @@ -1,3 +1,4 @@ + import logging from collections import OrderedDict @@ -6,6 +7,7 @@ import torch.nn as nn import trainer.networks as networks import trainer.lr_scheduler as lr_scheduler from .base_model import BaseModel +import torch_intermediary as ml logger = logging.getLogger('base') @@ -40,7 +42,8 @@ class FeatureModel(BaseModel): else: if self.rank <= 0: logger.warning('Params [{:s}] will not optimize.'.format(k)) - self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'], + # torch.optim.Adam + self.optimizer_G = ml.Adam(optim_params, lr=train_opt['lr_G'], weight_decay=wd_G, betas=(train_opt['beta1_G'], train_opt['beta2_G'])) self.optimizers.append(self.optimizer_G) diff --git a/codes/trainer/lr_scheduler.py b/codes/trainer/lr_scheduler.py index f437ff3c..ca867e6b 100644 --- a/codes/trainer/lr_scheduler.py +++ b/codes/trainer/lr_scheduler.py @@ -3,10 +3,10 @@ from collections import Counter from collections import defaultdict import torch from torch.optim.lr_scheduler import _LRScheduler +import torch_intermediary as ml from utils.util import opt_get - def get_scheduler_for_name(name, optimizers, scheduler_opt): schedulers = [] for o in optimizers: @@ -136,7 +136,8 @@ class CosineAnnealingLR_Restart(_LRScheduler): if __name__ == "__main__": - optimizer = torch.optim.Adam([torch.zeros(3, 64, 3, 3)], lr=1e-4, weight_decay=0, + #torch.optim.Adam + optimizer = ml.Adam([torch.zeros(3, 64, 3, 3)], lr=1e-4, weight_decay=0, betas=(0.9, 0.99)) ############################## # MultiStepLR_Restart diff --git a/codes/trainer/steps.py b/codes/trainer/steps.py index 9ec2f78b..2bf75617 100644 --- a/codes/trainer/steps.py +++ b/codes/trainer/steps.py @@ -12,6 +12,7 @@ from utils.util import recursively_detach, opt_get, clip_grad_norm logger = logging.getLogger('base') +import torch_intermediary as ml # Defines the expected API for a single training step class ConfigurableStep(Module): @@ -82,7 +83,8 @@ class ConfigurableStep(Module): import torch.nn as nn norm_modules = (nn.BatchNorm2d, nn.InstanceNorm2d, nn.BatchNorm1d, nn.InstanceNorm1d, nn.BatchNorm3d, nn.InstanceNorm3d, nn.GroupNorm, nn.LayerNorm) - emb_modules = (nn.Embedding, nn.EmbeddingBag) + # nn.Embedding + emb_modules = (ml.Embedding, nn.EmbeddingBag) param_names_notweights = set() all_param_names = set() param_map = {} @@ -123,7 +125,8 @@ class ConfigurableStep(Module): { 'params': params_weights, 'weight_decay': opt_get(opt_config, ['weight_decay'], 0) }, { 'params': params_notweights, 'weight_decay': 0 } ] - opt = torch.optim.AdamW(groups, lr=opt_config['lr'], + # torch.optim.AdamW + opt = ml.AdamW(groups, lr=opt_config['lr'], weight_decay=opt_get(opt_config, ['weight_decay'], 1e-2), betas=(opt_get(opt_config, ['beta1'], .9), opt_get(opt_config, ['beta2'], .999))) opt._group_names = [params_names_weights, params_names_notweights] @@ -141,14 +144,16 @@ class ConfigurableStep(Module): # The torch ZeRO implementation does not seem to support parameter groups, so do not shard the non-weighted # parameters and just use a normal AdamW implementation. In a large network, these weights will normally # be a tiny fraction of the total weights. - opt_unweighted = torch.optim.AdamW(params_notweights, lr=opt_config['lr'], weight_decay=0, + # torch.optim.AdamW + opt_unweighted = ml.AdamW(params_notweights, lr=opt_config['lr'], weight_decay=0, betas=(opt_get(opt_config, ['beta1'], .9), opt_get(opt_config, ['beta2'], .999))) opt_unweighted._config = opt_config opt_unweighted._config['network'] = net_name opt_unweighted._group_names = [] self.optimizers.append(opt_unweighted) - opt = ZeroRedundancyOptimizer(params_weights, optimizer_class=torch.optim.AdamW, lr=opt_config['lr'], + # torch.optim.AdamW + opt = ZeroRedundancyOptimizer(params_weights, optimizer_class=ml.AdamW, lr=opt_config['lr'], weight_decay=opt_get(opt_config, ['weight_decay'], 1e-2), betas=(opt_get(opt_config, ['beta1'], .9), opt_get(opt_config, ['beta2'], .999))) opt.param_groups[0]['initial_lr'] = opt_config['lr'] @@ -162,7 +167,8 @@ class ConfigurableStep(Module): opt._group_names = sorted(list(all_param_names)) elif self.step_opt['optimizer'] == 'lamb': from trainer.optimizers.lamb import Lamb - opt_unweighted = torch.optim.AdamW(params_notweights, lr=opt_config['lr'], weight_decay=0, + # torch.optim.AdamW + opt_unweighted = ml.AdamW(params_notweights, lr=opt_config['lr'], weight_decay=0, betas=(opt_get(opt_config, ['beta1'], .9), opt_get(opt_config, ['beta2'], .999))) opt_unweighted._config = opt_config opt_unweighted._config['network'] = net_name diff --git a/requirements.txt b/requirements.txt index 877d9869..3e3ebd3b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -45,4 +45,7 @@ rotary-embedding-torch axial_positional_embedding g-mlp-pytorch x-clip -x_transformers==1.0.4 \ No newline at end of file +x_transformers==1.0.4 + +# bitsandbytes +bitsandbytes==0.35.0 diff --git a/setup.py b/setup.py index 1073fb00..08c69087 100644 --- a/setup.py +++ b/setup.py @@ -67,6 +67,8 @@ setuptools.setup( "g-mlp-pytorch", "x-clip", "x_transformers==1.0.4", + + "bitsandbytes==0.35.0", ], classifiers=[ "Programming Language :: Python :: 3",