Блог

У меня есть несколько действительно больших входных тензоров, и при их сборке я столкнулся с проблемами с памятью, поэтому я читал их один за другим в .pt файл. Когда я запускаю скрипт, который генерирует и сохраняет файл, файл становится все больше и больше, поэтому я предполагаю, что тензоры сохраняются правильно. Вот этот код:

 with open(a_sync_save, "ab") as f:
     print("saved")
     torch.save(torch.unsqueeze(torch.cat(tensors, dim=0), dim=0), f)
  

Я хочу одновременно считывать определенное количество этих тензоров из файла, потому что я не хочу снова столкнуться с проблемой памяти.Когда я пытаюсь прочитать каждый тензор, сохраненный в файле, мне удается получить только первый тензор.

 with open(a_sync_save, "rb") as f:
    for tensor in torch.load(f):
        print(tensor.shape)
  

Вывод здесь представляет собой форму первого тензора, затем мирно завершается.

Вот некоторый код, который я использовал для ответа на этот вопрос. Многое из этого зависит от того, что я делаю, но суть этого может быть использована другими, которые сталкиваются с той же проблемой, что и я.

 def stream_training(filepath, epochs=100):
    """
    :param filepath: file path of pkl file
    :param epochs: number of epochs to run
    """
    def training(train_dataloader, model_obj, criterion, optimizer):
        for j, data in enumerate(train_dataloader, start=0):
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data
            inputs, labels = inputs.cuda(), labels.cuda()
            outputs = model_obj(inputs.float())
            outputs = torch.flatten(outputs)
            loss = criterion(outputs, labels.float())
            print(loss)
            # zero the parameter gradients
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model_obj.parameters(), max_norm=1)
            optimizer.step()

    tensors = []
    expected_values = []
    model= Model(1000, 1, 256, 1)
    model.cuda()
    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.00001, betas=(0.9, 0.99999), eps=1e-08, weight_decay=0.001,
                           amsgrad=True)
    for i in range(epochs):
        with (open(filepath, 'rb')) as openfile:
            while True:
                try:
                    data_list = pickle.load(openfile)
                    tensors.append(data_list[0])
                    expected_values.append(data_list[1])
                    if len(tensors) % BATCH_SIZE == 0:
                        tensors = torch.cat(tensors, dim=0)
                        tensors = torch.reshape(tensors, (tensors.shape[0], tensors.shape[1], -1))
                        train_loader = make_dataset(tensors, expected_values) # makes a dataloader for the batch that comes in
                        training(train_loader, model, criterion, optimizer)  #Performs forward and back prop
                        tensors = [] # washes out the batch to conserve memory on my computer.
                        expected_values = []
                except EOFError:
                    print("This file has finished training")
                    break
  

Вот модель для развлечения.

 class Model(nn.Module):
    def __init__(self, input_size, output_size, hidden_dim, n_layers):
        super(Model, self).__init__()
        # dimensions
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers

        #Define the layers
        #GRU
        self.gru = nn.GRU(input_size, hidden_dim, n_layers, batch_first=True)
        self.fc1 = nn.Linear(hidden_dim, hidden_dim)
        self.bn1 = nn.BatchNorm1d(num_features=hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.bn2 = nn.BatchNorm1d(num_features=hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, hidden_dim)
        self.bn3 = nn.BatchNorm1d(num_features=hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, hidden_dim)
        self.bn4 = nn.BatchNorm1d(num_features=hidden_dim)
        self.fc5 = nn.Linear(hidden_dim, hidden_dim)
        self.output = nn.Linear(hidden_dim, output_size)

    def forward(self, x):
        x = x.float()
        x = F.relu(self.gru(x)[1])
        x = x[-1,:,:] # eliminates first dim
        x = F.dropout(x, 0.5)
        x = F.relu(self.bn1(self.fc1(x)))
        x = F.dropout(x, 0.5)
        x = F.relu(self.bn2(self.fc2(x)))
        x = F.dropout(x, 0.5)
        x = F.relu(self.bn3(self.fc3(x)))
        x = F.dropout(x, 0.5)
        x = F.relu(self.bn4(self.fc4(x)))
        x = F.dropout(x, 0.5)
        x = F.relu(self.fc5(x))
        return torch.sigmoid(self.output(x))

    def init_hidden(self, batch_size):
        hidden = torch.zeros(self.n_layers, batch_size, self.hidden_dim)
        return hidden