Eu estava procurando maneiras alternativas de salvar um modelo treinado no PyTorch. Até agora, encontrei duas alternativas.

1. torch.save () para salvar um modelo e torch.load () para carregar um modelo.
2. model.state_dict () para salvar um modelo treinado e model.load_state_dict () para carregar o modelo salvo.

Eu me deparei com essa discussão em que a abordagem 2 é recomendada sobre a abordagem 1.

Minha pergunta é: por que a segunda abordagem é preferida? É apenas porque os módulos torch.nn têm essas duas funções e somos incentivados a usá-los?

Encontrei esta página no repositório do github, colarei o conteúdo aqui.

Abordagem recomendada para salvar um modelo

Existem duas abordagens principais para serializar e restaurar um modelo.

O primeiro (recomendado) salva e carrega apenas os parâmetros do modelo:

torch.save(the_model.state_dict(), PATH)

Depois, mais tarde:

the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

O segundo salva e carrega o modelo inteiro:

torch.save(the_model, PATH)

Depois, mais tarde:

the_model = torch.load(PATH)

No entanto, nesse caso, os dados serializados são vinculados às classes específicas e à estrutura de diretórios exata usada, para que possam ser quebrados de várias maneiras quando usados ​​em outros projetos ou após alguns refatores sérios.

Depende do que você quer fazer.

Caso nº 1: salve o modelo para usá-lo por inferência : salve o modelo, restaure-o e altere o modelo para o modo de avaliação. Isso é feito porque você geralmente tem BatchNorme Dropoutcamadas que por padrão estão no modo de trem na construção:

torch.save(model.state_dict(), filepath)

#Later to restore:
model.load_state_dict(torch.load(filepath))
model.eval()

Caso nº 2: salve o modelo para continuar o treinamento mais tarde : se você precisar continuar treinando o modelo que está prestes a salvar, precisará salvar mais do que apenas o modelo. Você também precisa salvar o estado do otimizador, épocas, pontuação etc. Você faria assim:

state = {
    'epoch': epoch,
    'state_dict': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    ...
}
torch.save(state, filepath)

Para retomar o treinamento, você faria coisas como: state = torch.load(filepath)e, em seguida, para restaurar o estado de cada objeto individual, algo como isto:

model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])

Como você está retomando o treinamento, NÃO ligue model.eval()depois de restaurar os estados ao carregar.

Caso nº 3: Modelo a ser usado por outra pessoa sem acesso ao seu código : No Tensorflow, você pode criar um .pbarquivo que define a arquitetura e os pesos do modelo. Isso é muito útil, especialmente ao usar Tensorflow serve. A maneira equivalente de fazer isso no Pytorch seria:

torch.save(model, filepath)

# Then later:
model = torch.load(filepath)

Dessa forma, ainda não é à prova de balas e, como o pytorch ainda está passando por muitas alterações, eu não o recomendaria.

A biblioteca pickle Python implementa protocolos binários para serializar e desserializar um objeto Python.

Quando você import torch(ou quando você usa PyTorch) que vai import picklepara você e você não precisa chamar pickle.dump()e pickle.load()diretamente, o que são os métodos para salvar e carregar o objeto.

Na verdade, torch.save()e torch.load()vai embrulhar pickle.dump()e pickle.load()para você.

A state_dictoutra resposta mencionada merece apenas mais algumas notas.

O state_dictque temos dentro do PyTorch? Na verdade, existem dois state_dicts.

O modelo PyTorch é torch.nn.Moduletem model.parameters()chamada para obter parâmetros learnable (w eb). Esses parâmetros aprendíveis, uma vez definidos aleatoriamente, serão atualizados ao longo do tempo à medida que aprendemos. Parâmetros aprendíveis são os primeiros state_dict.

O segundo state_dicté o ditado de estado do otimizador. Você se lembra que o otimizador é usado para melhorar nossos parâmetros aprendíveis. Mas o otimizador state_dicté fixo. Nada a aprender lá.

Como os state_dictobjetos são dicionários Python, eles podem ser facilmente salvos, atualizados, alterados e restaurados, adicionando uma grande modularidade aos modelos e otimizadores do PyTorch.

Vamos criar um modelo super simples para explicar isso:

import torch
import torch.optim as optim

model = torch.nn.Linear(5, 2)

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

print("Model weight:")    
print(model.weight)

print("Model bias:")    
print(model.bias)

print("---")
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

Este código produzirá o seguinte:

Model's state_dict:
weight   torch.Size([2, 5])
bias     torch.Size([2])
Model weight:
Parameter containing:
tensor([[ 0.1328,  0.1360,  0.1553, -0.1838, -0.0316],
        [ 0.0479,  0.1760,  0.1712,  0.2244,  0.1408]], requires_grad=True)
Model bias:
Parameter containing:
tensor([ 0.4112, -0.0733], requires_grad=True)
---
Optimizer's state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [140695321443856, 140695321443928]}]

Observe que este é um modelo mínimo. Você pode tentar adicionar uma pilha de

model = torch.nn.Sequential(
          torch.nn.Linear(D_in, H),
          torch.nn.Conv2d(A, B, C)
          torch.nn.Linear(H, D_out),
        )

Observe que apenas as camadas com parâmetros que podem ser aprendidos (camadas convolucionais, camadas lineares etc.) e buffers registrados (camadas batchnorm) têm entradas nos modelos state_dict.

Coisas não aprendidas pertencem ao objeto otimizador state_dict, que contém informações sobre o estado do otimizador, bem como os hiperparâmetros usados.

O resto da história é o mesmo; na fase de inferência (essa é uma fase em que usamos o modelo após o treinamento) para prever; nós previmos com base nos parâmetros que aprendemos. Portanto, para a inferência, só precisamos salvar os parâmetros model.state_dict().

torch.save(model.state_dict(), filepath)

E para usar mais tarde model.load_state_dict (torch.load (filepath)) model.eval ()

Nota: Não esqueça a última linha, model.eval()isto é crucial após o carregamento do modelo.

Também não tente salvar torch.save(model.parameters(), filepath). O model.parameters()é apenas o objeto gerador.

Por outro lado, torch.save(model, filepath)salva o próprio objeto do modelo, mas lembre-se de que o modelo não possui o otimizador state_dict. Verifique a outra excelente resposta de @Jadiel de Armas para salvar o ditado de estado do otimizador.

Uma convenção comum do PyTorch é salvar modelos usando uma extensão de arquivo .pt ou .pth.

Salvar / carregar modelo inteiro Salvar:

path = "username/directory/lstmmodelgpu.pth"
torch.save(trainer, path)

Carga:

A classe do modelo deve ser definida em algum lugar

model = torch.load(PATH)
model.eval()

Se você deseja salvar o modelo e deseja retomar o treinamento posteriormente:

GPU única: Salvar:

state = {
        'epoch': epoch,
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
}
savepath='checkpoint.t7'
torch.save(state,savepath)

Carga:

checkpoint = torch.load('checkpoint.t7')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint['epoch']

GPU múltipla: Salvar

state = {
        'epoch': epoch,
        'state_dict': model.module.state_dict(),
        'optimizer': optimizer.state_dict(),
}
savepath='checkpoint.t7'
torch.save(state,savepath)

Carga:

checkpoint = torch.load('checkpoint.t7')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint['epoch']

#Don't call DataParallel before loading the model otherwise you will get an error

model = nn.DataParallel(model) #ignore the line if you want to load on Single GPU