Como ler de um conjunto de dados de IO alto no pytorch que cresce de época para época

Eu uso o Tensorflow, mas estou escrevendo documentação para usuários que normalmente variam entre as estruturas de aprendizado profundo .

Ao trabalhar com conjuntos de dados que não se encaixam no sistema de arquivos local (TB +), colho dados de um repositório de dados remoto e gravo amostras localmente em um tfrecordsformato padrão do Tensorflow .

Durante a primeira época do treinamento, apenas amostrarei alguns valores; portanto, uma época de dados locais é muito pequena, eu treino nela. Na época 2 , re-examino quais arquivos de dados foram produzidos pelos meus subprocessos de amostragem (agora mais) e treino o conjunto expandido de arquivos de dados locais para a próxima época. Repita o processo a cada época. Dessa maneira, construo um cache local de amostras e posso despejar amostras mais antigas conforme preencho o armazenamento local. O cache de amostras locais cresce aproximadamente no momento em que o modelo precisa mais da variação (na última parte do treinamento).

No Python / Tensorflow, é crucial que eu não desserialize os dados no processo de loop de treinamento do Python, porque o GIL do Python não pode suportar as taxas de transferência de dados (300-600 MB / s, os dados são científicos não compactáveis) e, portanto, o desempenho da GPU sofre quando o GIL do Python não pode atender rapidamente ao ciclo de treinamento.

Gravar as amostras em um tfrecordsarquivo a partir de subprocessos (multiprocessamento python) permite que o nativo do tensorflow TFRecordsDatasetfaça a desserialização fora do Python e, assim, contornamos os problemas do GIL do Python, e posso saturar uma GPU com altas taxas de dados de IO.

Gostaria de saber como resolveria esse problema em Pytorch. Estou escrevendo sobre a estratégia de amostragem que está sendo usada e quero fornecer recomendações específicas aos usuários do Tensorflow e do PyTorch, mas não conheço o ecossistema de pré-processamento do PyTorch o suficiente para escrever com detalhes suficientes.

Nota lateral: a única solução puramente baseada em Python para suportar essas taxas de transferência de dados pode vir no Python 3.8 com memória compartilhada e multiprocessamento do System V, mas ainda não tentei isso, pois o suporte a ela não é suficiente (em breve será ) As soluções de multiprocessamento existentes não são suficientes porque exigem desserialização no processo do loop de treinamento e, portanto, bloqueiam o GIL durante a desserialização a altas taxas de IO.

Na verdade, você pode facilmente desserializar dados em um subprocesso usando torch.utils.data.DataLoader. Ao definir o num_workersargumento como 1 ou um valor maior, você pode gerar subprocessos com seus próprios intérpretes de python e GILs.

loader = torch.utils.data.DataLoader(your_dataset, num_workers=n, **kwargs)
for epoch in range(epochs):
    for batch_idx, data in enumerate(loader):
         # loader in the main process does not claim GIL at this point

A Dataloaderrequer a torch.utils.data.Datasetpara obter dados. Pode não ser um trabalho trivial implementar uma subclasse adequada no seu caso. Caso seja necessário recriar uma Datasetinstância para cada época, você pode fazer algo assim.

for epcoh in range(epochs):
    dset = get_new_dataset()
    loader = torch.utils.data.DataLoader(dset, num_workers=n, **kwargs)
    for batch_idx, data in enumerate(loader):
        # Do training

ou melhor ainda

dset = get_new_dataset()
loader = torch.utils.data.DataLoader(dset, num_workers=n, **kwargs)

for epcoh in range(epochs):
    last_batch_idx =  (len(dset)-1) // loader.batch_size
    for batch_idx, data in enumerate(loader):
        # Prepare next loader in advance to avoid blocking
        if batch_idx == last_batch_idx:
            dset = get_new_dataset()
            loader = torch.utils.data.DataLoader(dset, num_workers=n, **kwargs)
        # Do training

Como uma observação lateral, observe que a operação vinculada à CPU é afetada pelo GIL na maioria dos casos, não a operação vinculada a E / S, ou seja, threadingservirá para qualquer operação pesada de E / S e você nem precisa subprocess. Para mais informações, consulte esta pergunta e este artigo da Wikipedia .