CPU memory increase while training

[gudrb]

During the training step, the code block:

for blk in self.v.blocks: x = blk(x) causes an increase in CPU memory, eventually leading to full CPU memory. Do you know any solutions to this issue?

[YuanGongND]

hi there,

Are you using our recipe to train the model?

are you training on cpu or gpu? What is you cpu/gpu memory?

x = blk(x) means the input tensor passes a Transformer layer, it should take GPU memory. I don't know how torch manages cpu memory in this case.

-Yuan

[gudrb]

i am using my own training code but following the ast_models.py code to define ASTModel. i can see from the memory_profiler, at line 176 of screenshot, my htop cpu memory linearly increase while trianing step.

--> This is part of my code.

class ASTModel(nn.Module): """ The AST model. :param label_dim: the label dimension, i.e., the number of total classes, it is 527 for AudioSet, 50 for ESC-50, and 35 for speechcommands v2-35 :param fstride: the stride of patch spliting on the frequency dimension, for 1616 patchs, fstride=16 means no overlap, fstride=10 means overlap of 6 :param tstride: the stride of patch spliting on the time dimension, for 1616 patchs, tstride=16 means no overlap, tstride=10 means overlap of 6 :param input_fdim: the number of frequency bins of the input spectrogram :param input_tdim: the number of time frames of the input spectrogram :param imagenet_pretrain: if use ImageNet pretrained model :param audioset_pretrain: if use full AudioSet and ImageNet pretrained model :param model_size: the model size of AST, should be in [tiny224, small224, base224, base384], base224 and base 384 are same model, but are trained differently during ImageNet pretraining. """ def init(self, label_dim=527, input_fdim=128, input_tdim=1024, imagenet_pretrain=True, audioset_pretrain=False, model_size='base384', verbose=True):

    super(ASTModel, self).__init__()
    assert timm.__version__ == '0.4.5', 'Please use timm == 0.4.5, the code might not be compatible with newer versions.'

    if verbose == True:
        print('---------------AST Model Summary---------------')
        print('ImageNet pretraining: {:s}, AudioSet pretraining: {:s}'.format(str(imagenet_pretrain),str(audioset_pretrain)))
    # override timm input shape restriction
    # timm.models.vision_transformer.PatchEmbed = PatchEmbed
    # if AudioSet pretraining is not used (but ImageNet pretraining may still apply)
    if audioset_pretrain == False:
        if model_size == 'tiny224':
            self.v = timm.create_model('vit_deit_tiny_distilled_patch16_224', pretrained=imagenet_pretrain)
        elif model_size == 'small224':
            self.v = timm.create_model('vit_deit_small_distilled_patch16_224', pretrained=imagenet_pretrain)
        elif model_size == 'base224':
            self.v = timm.create_model('vit_deit_base_distilled_patch16_224', pretrained=imagenet_pretrain)
        elif model_size == 'base384':
            self.v = timm.create_model('vit_deit_base_distilled_patch16_384', pretrained=imagenet_pretrain)
        else:
            raise Exception('Model size must be one of tiny224, small224, base224, base384.')

def get_shape(self, fstride, tstride, input_fdim=128, input_tdim=1024):
    test_input = torch.randn(1, 1, input_fdim, input_tdim)
    test_proj = nn.Conv2d(1, self.original_embedding_dim, kernel_size=(16, 16), stride=(fstride, tstride))
    test_out = test_proj(test_input)
    f_dim = test_out.shape[2]
    t_dim = test_out.shape[3]
    return f_dim, t_dim

@autocast()
def forward(self, x):
    """
    :param x: the input spectrogram, expected shape: (batch_size, time_frame_num, frequency_bins), e.g., (12, 1024, 128)
    :return: prediction
    """
    # expect input x = (batch_size, time_frame_num, frequency_bins), e.g., (12, 1024, 128)
    # x = x.unsqueeze(1)
    # x = x.transpose(1, 2)
    B,T,F = x.shape
    # x = self.v.patch_embed(x)
    cls_tokens = self.v.cls_token.expand(B, -1, -1)
    dist_token = self.v.dist_token.expand(B, -1, -1)
    x = torch.cat((cls_tokens, dist_token, x), dim=1)
    x = x + self.v.pos_embed
    x = self.v.pos_drop(x)
    for blk in self.v.blocks:
        x = blk(x)
    x = self.v.norm(x)

---

--> i am only utilizing pretrained raw vit_deit_tiny_distilled_patch16_224 without changing patch embedding or positional encoding.

class LSQSL(nn.Module):

def init(self, in_chans=1, frequency=129, d1_model=32, d2_model=172, n_layers=1, hidden_size=88, bidirectional=True):

def __init__(self, in_chans=1, time=29, frequency=129, n_seq=10, seq_length=20, d1_model=32, d2_model=128, n_layers=1, hidden_size=64, bidirectional=True, head_num=2):
    super(LSQSL, self).__init__()
    self.ast_mdl = ASTModel(label_dim=5, \
     input_fdim=192, input_tdim=196, \
     imagenet_pretrain=True, audioset_pretrain=False, \
     model_size='tiny224')

def forward(self, raw):
    B, L, D = raw.shape
    raw = self.ast_mdl(raw)
    return raw

---

--> i use GPU for training with (model = LSQSL.to(device)). but the blk(x) increase my htop cpu memory linearly increase. when i track the code until /python3.7/site-packages/timm/models/vision_transformer.py, in the Block class, self.attn and self.mlp operation increase the cpu memory.

def forward(self, x): x = x + self.drop_path(self.attn(self.norm1(x))) x = x + self.drop_path(self.mlp(self.norm2(x))) return x

class Mlp(nn.Module): def init(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): super().init() out_features = out_features or in_features hidden_features = hidden_features or in_features self.fc1 = nn.Linear(in_features, hidden_features) self.act = act_layer() self.fc2 = nn.Linear(hidden_features, out_features) self.drop = nn.Dropout(drop) def forward(self, x): x = self.fc1(x) x = self.act(x) x = self.drop(x) x = self.fc2(x) x = self.drop(x) return x

class Attention(nn.Module): def init(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): super().init() self.num_heads = num_heads head_dim = dim // num_heads

NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights

    self.scale = qk_scale or head_dim ** -0.5
    self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
    self.attn_drop = nn.Dropout(attn_drop)
    self.proj = nn.Linear(dim, dim)
    self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x):
    B, N, C = x.shape
    qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
    q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)

    attn = (q @ k.transpose(-2, -1)) * self.scale
    attn = attn.softmax(dim=-1)
    attn = self.attn_drop(attn)

    x = (attn @ v).transpose(1, 2).reshape(B, N, C)
    x = self.proj(x)
    x = self.proj_drop(x)
    return x

--> However, Mlp and Attention class foward have nothing related about cpu operation or stack of trash variables such as detach() or append(). Do u know the reason of this unexpected CPU memory leack? i am only suspecting the timm library own Issues.

[YuanGongND]

What's your GPU memory and CPU memory?

Have you tried to use our recipe (e.g., ESC-50, which is very fast to run, see our Readme) and see if the issue is still there?

This is the first time I heard this issue.

[gudrb]

I also see this memory leak from your ESC-50 training code also. but, not as much as mine. I think the difference from mine is sequence length. I use original 198 sequence length with much more attention computation.

I am using GPU: NVIDIA A100-PCIE-40GB, and CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 46 bits physical, 48 bits virtual Byte Order: Little Endian CPU(s): 48 On-line CPU(s) list: 0-47 Vendor ID: GenuineIntel Model name: Intel Xeon Processor (Skylake, IBRS) CPU family: 6 Model: 85 Thread(s) per core: 1 Core(s) per socket: 1 Socket(s): 48 Stepping: 4 BogoMIPS: 6185.46 Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology cpuid tsc_k nown_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 arat pku ospke avx512_vnni md_clear Virtualization features: Hypervisor vendor: KVM Virtualization type: full Caches (sum of all):
L1d: 1.5 MiB (48 instances) L1i: 1.5 MiB (48 instances) L2: 192 MiB (48 instances) L3: 768 MiB (48 instances) NUMA:
NUMA node(s): 1 NUMA node0 CPU(s): 0-47 Vulnerabilities:
Gather data sampling: Unknown: Dependent on hypervisor status Itlb multihit: KVM: Mitigation: VMX unsupported L1tf: Mitigation; PTE Inversion Mds: Mitigation; Clear CPU buffers; SMT Host state unknown Meltdown: Mitigation; PTI Mmio stale data: Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown Retbleed: Mitigation; IBRS Spec rstack overflow: Not affected Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization Spectre v2: Mitigation; IBRS, IBPB conditional, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected Srbds: Not affected Tsx async abort: Mitigation; Clear CPU buffers; SMT Host state unknown

[gudrb]

Thank you for answering my questions. It was not any problem of the timm library or the AST codes. While training, there was a tensor operation without using .item() such <running_loss += loss> So i changed it to <running_loss += loss.item()> and RAM memory leak was fixed!

[YuanGongND]

thanks for letting me know. Good luck with your research!