How to prevent memory use growth when updating weights and biases in a Pytorch model

user_14366234 0
Viewed 208

I'm trying to build a VGG16 model to make an ONNX export using Pytorch. I want to force the model with my own set of weights and biases. But in this process my computer quickly runs out of memory.

Here is how I want to do it (this is only a test, in the real version I read the weights and biases in a set of files), this example only force all values to 0.5

# Create empty VGG16 model (random weights)
from torchvision import models
from torchsummary import summary

vgg16 = models.vgg16()
# la structure est : vgg16.__dict__
summary(vgg16, (3, 224, 224))

#  convolutive layers
for layer in vgg16.features:
    print()
    print(layer)
    if (hasattr(layer,'weight')):
        dim = layer.weight.shape
        print(dim)
        print(str(dim[0]*(dim[1]*dim[2]*dim[3]+1))+' params')

        # Remplacement des poids et biais
        for i in range (dim[0]):
            layer.bias[i] = 0.5
            for j in range (dim[1]):
                for k in range (dim[2]):
                    for l in range (dim[3]):
                        layer.weight[i][j][k][l] = 0.5

# Dense layers
for layer in vgg16.classifier:
    print()
    print(layer)
    if (hasattr(layer,'weight')):
        dim = layer.weight.shape
        print(str(dim)+' --> '+str(dim[0]*(dim[1]+1))+' params')
        for i in range(dim[0]):
            layer.bias[i] = 0.5
            for j in range(dim[1]):
                layer.weight[i][j] = 0.5

When I look at the memory usage of the computer, it grows linealrly and saturates the 16GB RAM during the first dense layer processing. Then python crashes...

Is there another better way to do this, keeping in mind that I want to onnx export the model afterwards? Thanks for your help.

pytorch deep-learning machine-learning neural-network python
4 Answers

The memory growth is caused by the need to adjust gradient for every weight and bias change. Try setting .requires_grad attribute to False before the update and restoring it after the update. Example:

for layer in vgg16.features:
    print()
    print(layer)
    if (hasattr(layer,'weight')):
        
        # supress .requires_grad
        layer.bias.requires_grad = False
        layer.weight.requires_grad = False
        
        dim = layer.weight.shape
        print(dim)
        print(str(dim[0]*(dim[1]*dim[2]*dim[3]+1))+' params')

        # Remplacement des poids et biais
        for i in range (dim[0]):
            layer.bias[i] = 0.5
            for j in range (dim[1]):
                for k in range (dim[2]):
                    for l in range (dim[3]):
                        layer.weight[i][j][k][l] = 0.5
        
        # restore .requires_grad
        layer.bias.requires_grad = True
        layer.weight.requires_grad = True
user_103964690

I also found here a solution based on creating a dictionary with the new weights and using it to update the model:

weight_dict.update(new_weight_dict)

What do you think about it? Is there any documentation about this method? Maybe it's faster?

user_143662340

I finally did it like this : it does not saturate the RAM and is much much faster than my previous version.

# Create empty VGG16 model (random weights)
import torch
from torchvision import models
from torchsummary import summary

vgg16 = models.vgg16()
# The structure is in : vgg16.__dict__
summary(vgg16, (3, 224, 224))
state_dict = vgg16.state_dict()
print(state_dict.keys())

# Prepare to read files
import os
workdir = 'C:\\Users\\...\\'
path = workdir + '...\\'

# Convolutive layers
index = 1
nCouche = 0
with torch.no_grad():
    for layer in vgg16.features:
        print()
        print(layer)
        if (hasattr(layer,'weight')):

            dim = layer.weight.shape
            print(str(dim) + ' : ' +str(dim[0]*(dim[1]*dim[2]*dim[3]+1))+' parameters')

            # Read the weights
            nomFichier = nomCouche + '_weights.syntxt'
            print('Reading file: ' + nomFichier)
            file = open (path + nomFichier,"r")
            weights = []
            for line in file:
                wprov = [int(x) for x in line.split()]
                weights.append([[[wprov[(k*dim[3]+j)*dim[2]+i] for i in range(dim[3])] for j in range(dim[2])] for k in range(dim[1])])
            file.close()

            # Read the biases
            nomFichier = nomCouche + '_biases.syntxt'
            print('Reading file: ' + nomFichier)
            file = open (path + nomFichier,"r")
            biases = []
            for line in file:
                biases.append(int(line))
            file.close()

            # Replace weights and biases
            nameW = 'features.'+str(nCouche)+'.weight'
            nameB = 'features.'+str(nCouche)+'.bias'
            print ('Features: ' + nameW + ' & ' + nameB)
            w = torch.empty(dim)
            b = torch.empty(dim[0])

            w = torch.IntTensor(weights)
            b = torch.IntTensor(biases)
            state_dict[nameW].copy_(w)
            state_dict[nameB].copy_(b)

            index += 1
        nCouche += 1

    index = 1
    nCouche = 0
    # Dense layers
    for layer in vgg16.classifier:
        print()
        print(layer)
        if (hasattr(layer,'weight')):
            dim = layer.weight.shape
            print(str(dim)+' --> '+str(dim[0]*(dim[1]+1))+' params')

            # Prepare file names
            nomFichier = nomCouche + '_weights.syntxt'
            print('Reading file: ' + nomFichier)
            file = open (path+nomFichier,"r")

            # Prepare to replace weights and biases
            nameW = 'classifier.'+str(nCouche)+'.weight'
            nameB = 'classifier.'+str(nCouche)+'.bias'
            print ('Features: ' + nameW + ' & ' + nameB)
            w = torch.empty(dim)
            b = torch.empty(dim[0])

            # Read the weights
            i = 0
            interval = 500
            print('Countdown:') # because it can be quite long
            for line in file:
                weights = [int(x) for x in line.split()]
                w[i,:] = torch.IntTensor(weights)
                i += 1
                if ((dim[0] - i) % interval == 0):
                    print(f"{i/dim[0]*100:.1f}"+'%')
            print('Read ' + str(i) + ' lines of ' + str(len(weights)) + ' weights')
            file.close()

            # Read the biases
            nomFichier = 'dense'+str(index)+'_biases.syntxt'
            print('Reading file: ' + nomFichier)
            file = open (path+nomFichier,"r")
            i = 0
            for line in file:
                biases = [int(y) for y in line.split()]
                b[i] = biases[0]
                i += 1
            file.close()

            state_dict[nameW].copy_(w)
            state_dict[nameB].copy_(b)

            index += 1
        nCouche += 1

# Update state dictionary
vgg16.load_state_dict(state_dict)

It reads the weights and biases in a series of files, and replace them in the various tensors of the net. Maybe it can help someone who may have a similar problem...

user_143662340

If you can provide new weights as complete tensors - there is no need to update them one by one in nested cycles:

for layer in vgg16.features:
    if (hasattr(layer,'weight')):
        # Remplacement des poids et biais
        new_bias = torch.rand (layer.bias.shape)  # <-- YOUR DATA GOES HERE
        new_weight = torch.rand (layer.weight.shape) # <-- YOUR DATA GOES HERE
        layer_bias = torch.nn.Parameter(new_bias, requires_grad=True) 
        layer.weight = torch.nn.Parameter(new_weight, requires_grad=True)
user_103964690
Related