CUDA не хватает памяти при настройке большой модели - PullRequest
Новая
       19

CUDA не хватает памяти при настройке большой модели

0 голосов
/ 22 января 2019

Ранее я обучал режиму VGG (скажем, модель1) и двухслойной модели (скажем, модель2) отдельно, теперь мне нужно обучить новую модель, которая объединяет эти две модели вместе, и каждая часть новой модели инициализируетсяс изученными весами model1 и model2, которые я реализовал следующим образом: 

class TransferModel(nn.Module):
    def __init__(self, VGG, TwoLayer):
        super(TransferModel, self).__init__()
        self.vgg_layer=VGG
        self.linear = TwoLayer
        for param in self.vgg_layer.parameters():
            param.requires_grad = True
    def forward(self, x):
        h1_vgg = self.vgg_layer(x)
        y_pred = self.linear(h1_vgg)
        return y_pred
# for image_id in train_ids[0:1]:
#     img = load_image(train_id_to_file[image_id])
new_model=TransferModel(trained_vgg_instance, trained_twolayer_instance)
new_model.linear.load_state_dict(trained_twolayer_instance.state_dict())
new_model.vgg_layer.load_state_dict(trained_vgg_instance.state_dict())
new_model.cuda()

А при обучении я пытаюсь: 

def train(model, learning_rate=0.001, batch_size=50, epochs=2):
    optimizer=optim.Adam(model.parameters(), lr=learning_rate)
    criterion = torch.nn.MultiLabelSoftMarginLoss()
    x = torch.zeros([batch_size, 3, img_size, img_size])
    y_true = torch.zeros([batch_size, 4096])
    for epoch in range(epochs):  # loop over the dataset multiple times
        running_loss = 0.0
        shuffled_indcs=torch.randperm(20000)
        for i in range(20000):
        for batch_num in range(int(20000/batch_size)):
            optimizer.zero_grad()
            for j in range(batch_size):
                # ... some code to load batches of images into x....
            x_batch=Variable(x).cuda()
            print(batch_num)
            y_true_batch=Variable(train_labels[batch_num*batch_size:(batch_num+1)*batch_size, :]).cuda()
            y_pred =model(x_batch)
            loss = criterion(y_pred, y_true_batch)
            loss.backward()
            optimizer.step()
            running_loss += loss
            del x_batch, y_true_batch, y_pred
            torch.cuda.empty_cache()
        print("in epoch[%d] = %.8f " % (epoch, running_loss /(batch_num+1)))
        running_loss = 0.0

    print('Finished Training')
train(new_model)

Во второй итерации (batch_num = 1) изПервая эпоха, я получаю эту ошибку:

CUDA из памяти.Попытка выделить 153,12 МБ (ГП 0; 5,93 ГБ общей емкости; 4,83 ГБ уже выделено; 66,94 МБ свободной; 374,12 МБ кэшировано)

Хотя я явно использовал «del» в своей тренировке, выполнивnvidia-smi похоже, что ничего не делает и память не освобождается.

Что мне делать?

1 Ответ

0 голосов
/ 23 января 2019

Измените эту строку:

running_loss += loss

на это:

running_loss += loss.item()

Путем добавления loss к running_loss,Вы говорите Pytorch, чтобы сохранить все градиенты относительно loss для этого пакета в памяти, даже когда вы начинаете тренироваться на следующем пакете.Pytorch считает, что, возможно, вам захочется использовать running_loss в какой-либо функции больших потерь для нескольких пакетов позже, и поэтому сохраняет все градиенты (и, следовательно, активации) для всех пакетов в памяти.

Добавляя .item()Вы просто получаете потерю как питон float, а не torch.FloatTensor.Этот float отделен от графа pytorch и, таким образом, pytorch знает, что вам не нужны градиенты по отношению к нему.

Если вы используете старую версию pytorch без .item(), вы можете попробовать:

running_loss += float(loss).cpu().detach

Это также может быть вызвано аналогичной ошибкой в ​​цикле test(), если она есть.

...