Pytorch - Guardando y Exportando modelos

En posts anteriores hemos aprendido a utilizar la librería Pytorch, viendo los conceptos baśicos, cómo diseñar y entrenar redes neuroanles y a manejar datasets de manera eficiente. Sin embargo, entrenar un modelo es solo parte del trabajo. Una vez tenemos nuestra red lista necesitamos poder guardarla en un archivo, o exportarla, para luego importarla en nuestras aplicaciones y ponerla a trabajar en un entorno de producción. En este post vamos a ver las diferentes opciones que Pytorch nos ofrece a a la hora de exportar modelos.

Guardando modelos

Guardando los parámetros

La primera opción consiste en guardar sólo los parámetros de la red. Para ello, Pytorch nos permite guardar el state_dict del modelo, un dict de Python que contiene una relación directa entre todas las capas con parámetros de la red y sus valores.

En el siguiente código, utilizado ya en este post, entrenamos un modelo simple con el dataset MNIST y, una vez entrenado, guardamos el state_dict del modelo.

import torch
import torchvision
import torchvision.transforms as transforms
from tqdm import tqdm
import numpy as np

device = "cuda" if torch.cuda.is_available() else "cpu"
device

'cuda'

# preparamos los datos

dataloader = {
    'train': torch.utils.data.DataLoader(torchvision.datasets.MNIST('../data', train=True, download=True,
                       transform=torchvision.transforms.Compose([
                            torchvision.transforms.ToTensor(),
                            torchvision.transforms.Normalize((0.1307,), (0.3081,))
                            ])
                      ), batch_size=2048, shuffle=True, pin_memory=True),
    'test': torch.utils.data.DataLoader(torchvision.datasets.MNIST('../data', train=False,
                   transform=torchvision.transforms.Compose([
                        torchvision.transforms.ToTensor(),
                        torchvision.transforms.Normalize((0.1307,), (0.3081,))
                        ])
                     ), batch_size=2048, shuffle=False, pin_memory=True)
}

# definimos el modelo

def block(c_in, c_out, k=3, p=1, s=1, pk=2, ps=2):
    return torch.nn.Sequential(
        torch.nn.Conv2d(c_in, c_out, k, padding=p, stride=s),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(pk, stride=ps)
    )

class CNN(torch.nn.Module):
  def __init__(self, n_channels=1, n_outputs=10):
    super().__init__()
    self.conv1 = block(n_channels, 64)
    self.conv2 = block(64, 128)
    self.fc = torch.nn.Linear(128*7*7, n_outputs)

  def forward(self, x):
    x = self.conv1(x)
    x = self.conv2(x)
    x = x.view(x.shape[0], -1)
    x = self.fc(x)
    return x

# entrenamos el modelo

def fit(model, dataloader, epochs=5):
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = torch.nn.CrossEntropyLoss()
    for epoch in range(1, epochs+1):
        model.train()
        train_loss, train_acc = [], []
        bar = tqdm(dataloader['train'])
        for batch in bar:
            X, y = batch
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()
            y_hat = model(X)
            loss = criterion(y_hat, y)
            loss.backward()
            optimizer.step()
            train_loss.append(loss.item())
            acc = (y == torch.argmax(y_hat, axis=1)).sum().item() / len(y)
            train_acc.append(acc)
            bar.set_description(f"loss {np.mean(train_loss):.5f} acc {np.mean(train_acc):.5f}")
        bar = tqdm(dataloader['test'])
        val_loss, val_acc = [], []
        model.eval()
        with torch.no_grad():
            for batch in bar:
                X, y = batch
                X, y = X.to(device), y.to(device)
                y_hat = model(X)
                loss = criterion(y_hat, y)
                val_loss.append(loss.item())
                acc = (y == torch.argmax(y_hat, axis=1)).sum().item() / len(y)
                val_acc.append(acc)
                bar.set_description(f"val_loss {np.mean(val_loss):.5f} val_acc {np.mean(val_acc):.5f}")
        print(f"Epoch {epoch}/{epochs} loss {np.mean(train_loss):.5f} val_loss {np.mean(val_loss):.5f} acc {np.mean(train_acc):.5f} val_acc {np.mean(val_acc):.5f}")

model = CNN()
fit(model, dataloader)

loss 0.60396 acc 0.82169: 100%|██████████| 30/30 [00:08<00:00,  3.62it/s]
val_loss 0.18738 val_acc 0.94626: 100%|██████████| 5/5 [00:01<00:00,  4.34it/s]
  0%|          | 0/30 [00:00<?, ?it/s]

Epoch 1/5 loss 0.60396 val_loss 0.18738 acc 0.82169 val_acc 0.94626


loss 0.13297 acc 0.96153: 100%|██████████| 30/30 [00:07<00:00,  3.93it/s]
val_loss 0.08688 val_acc 0.97479: 100%|██████████| 5/5 [00:01<00:00,  4.42it/s]
  0%|          | 0/30 [00:00<?, ?it/s]

Epoch 2/5 loss 0.13297 val_loss 0.08688 acc 0.96153 val_acc 0.97479


loss 0.08036 acc 0.97647: 100%|██████████| 30/30 [00:07<00:00,  3.96it/s]
val_loss 0.06439 val_acc 0.98044: 100%|██████████| 5/5 [00:01<00:00,  4.45it/s]
  0%|          | 0/30 [00:00<?, ?it/s]

Epoch 3/5 loss 0.08036 val_loss 0.06439 acc 0.97647 val_acc 0.98044


loss 0.05911 acc 0.98280: 100%|██████████| 30/30 [00:07<00:00,  3.96it/s]
val_loss 0.04986 val_acc 0.98458: 100%|██████████| 5/5 [00:01<00:00,  4.36it/s]
  0%|          | 0/30 [00:00<?, ?it/s]

Epoch 4/5 loss 0.05911 val_loss 0.04986 acc 0.98280 val_acc 0.98458


loss 0.04988 acc 0.98510: 100%|██████████| 30/30 [00:07<00:00,  3.81it/s]
val_loss 0.04321 val_acc 0.98639: 100%|██████████| 5/5 [00:01<00:00,  4.38it/s]

Epoch 5/5 loss 0.04988 val_loss 0.04321 acc 0.98510 val_acc 0.98639

# guardar modelo

PATH = './checkpoint.pt'
torch.save(model.state_dict(), PATH)

Ahora podemos cargar nuestro modelo y utilizarlo normalmente

# cargar modelo

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

CNN(
  (conv1): Sequential(
    (0): Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (conv2): Sequential(
    (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (fc): Linear(in_features=6272, out_features=10, bias=True)
)

def evaluate(model, dataloader):
    model.eval()
    model.to(device)
    bar = tqdm(dataloader['test'])
    acc = []
    with torch.no_grad():
        for batch in bar:
            X, y = batch
            X, y = X.to(device), y.to(device)
            y_hat = model(X)
            acc.append((y == torch.argmax(y_hat, axis=1)).sum().item() / len(y))
            bar.set_description(f"acc {np.mean(acc):.5f}")

evaluate(model, dataloader)

acc 0.98639: 100%|██████████| 5/5 [00:01<00:00,  4.34it/s]

Si bien de esta manera podemos guardar y cargar el modelo de manera eficiente, necesitamos tener un modelo instanciado para poder llamar a la función model.load_state_dict(). Esto significa que necesitaremos la definición de nuestro modelo allá dónde queramos importarlo (lo cual es poco flexible). Alternativamente, Pytorch nos permite guardar el modelo entero, y no solo el state_dict, de la siguiente manera.

torch.save(model, 'model.pt')

Y podemos cargar y evaluar nuestro modelo así

model = torch.load('model.pt')
model.eval()

CNN(
  (conv1): Sequential(
    (0): Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (conv2): Sequential(
    (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (fc): Linear(in_features=6272, out_features=10, bias=True)
)

evaluate(model, dataloader)

acc 0.98524: 100%|██████████| 5/5 [00:01<00:00,  4.69it/s]

Si bien de esta forma no necesitamos que nuestro modelo esté instanciado, seguimos necesitando su definición. Es por este motivo que la opción anterior es la recomendada, ya que es más eficiente (sólo guardamos los pesos) y también más flexible (podemos guardar otra información además del state_dict de nuestro modelo). Esta opción es ideal para guardar y cargar modelos durante el entrenamiento del mismo, quizás incluso junto al estado del optimizador, de manera que podemos entrenar modelos a partir de estos checkpoints en lugar de empezar de cero cada vez. Otro ejemplo consistiría en guardar el state_dict del modelo durante el entrenamiento solo cuando mejore una métrica determinada y cargar el mejor modelo al final del entrenamiento (que no tiene porqué coincidir con el último).

⚡ Aprende más sobre el guardado de modelos en Pytorch aquí.

def fit(model, dataloader, epochs=5, PATH="./checkpoint.pt"):
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = torch.nn.CrossEntropyLoss()
    best_acc = 0
    for epoch in range(1, epochs+1):
        model.train()
        train_loss, train_acc = [], []
        bar = tqdm(dataloader['train'])
        for batch in bar:
            X, y = batch
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()
            y_hat = model(X)
            loss = criterion(y_hat, y)
            loss.backward()
            optimizer.step()
            train_loss.append(loss.item())
            acc = (y == torch.argmax(y_hat, axis=1)).sum().item() / len(y)
            train_acc.append(acc)
            bar.set_description(f"loss {np.mean(train_loss):.5f} acc {np.mean(train_acc):.5f}")
        bar = tqdm(dataloader['test'])
        val_loss, val_acc = [], []
        model.eval()
        with torch.no_grad():
            for batch in bar:
                X, y = batch
                X, y = X.to(device), y.to(device)
                y_hat = model(X)
                loss = criterion(y_hat, y)
                val_loss.append(loss.item())
                acc = (y == torch.argmax(y_hat, axis=1)).sum().item() / len(y)
                val_acc.append(acc)
                bar.set_description(f"val_loss {np.mean(val_loss):.5f} val_acc {np.mean(val_acc):.5f}")
        # guardar modelo si es el mejor
        val_acc = np.mean(val_acc)
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), PATH)
            print(f"Best model saved at epoch {epoch} with val_acc {val_acc:.5f}")
        print(f"Epoch {epoch}/{epochs} loss {np.mean(train_loss):.5f} val_loss {np.mean(val_loss):.5f} acc {np.mean(train_acc):.5f} val_acc {np.mean(val_acc):.5f}")
    # cargar el mejor modelo al final del entrenamiento
    model.load_state_dict(torch.load(PATH))

model = CNN()
fit(model, dataloader)

loss 0.58207 acc 0.82621: 100%|██████████| 30/30 [00:07<00:00,  3.80it/s]
val_loss 0.17891 val_acc 0.94974: 100%|██████████| 5/5 [00:01<00:00,  4.38it/s]
  0%|          | 0/30 [00:00<?, ?it/s]

Best model saved at epoch 1 with val_acc 0.94974
Epoch 1/5 loss 0.58207 val_loss 0.17891 acc 0.82621 val_acc 0.94974


loss 0.13597 acc 0.96013: 100%|██████████| 30/30 [00:07<00:00,  3.96it/s]
val_loss 0.08804 val_acc 0.97444: 100%|██████████| 5/5 [00:01<00:00,  4.36it/s]
  0%|          | 0/30 [00:00<?, ?it/s]

Best model saved at epoch 2 with val_acc 0.97444
Epoch 2/5 loss 0.13597 val_loss 0.08804 acc 0.96013 val_acc 0.97444


loss 0.08168 acc 0.97642: 100%|██████████| 30/30 [00:07<00:00,  3.82it/s]
val_loss 0.06352 val_acc 0.98124: 100%|██████████| 5/5 [00:01<00:00,  4.37it/s]
  0%|          | 0/30 [00:00<?, ?it/s]

Best model saved at epoch 3 with val_acc 0.98124
Epoch 3/5 loss 0.08168 val_loss 0.06352 acc 0.97642 val_acc 0.98124


loss 0.06125 acc 0.98203: 100%|██████████| 30/30 [00:07<00:00,  3.82it/s]
val_loss 0.04724 val_acc 0.98474: 100%|██████████| 5/5 [00:01<00:00,  4.33it/s]
  0%|          | 0/30 [00:00<?, ?it/s]

Best model saved at epoch 4 with val_acc 0.98474
Epoch 4/5 loss 0.06125 val_loss 0.04724 acc 0.98203 val_acc 0.98474


loss 0.05294 acc 0.98442: 100%|██████████| 30/30 [00:07<00:00,  3.85it/s]
val_loss 0.04511 val_acc 0.98524: 100%|██████████| 5/5 [00:01<00:00,  4.14it/s]

Best model saved at epoch 5 with val_acc 0.98524
Epoch 5/5 loss 0.05294 val_loss 0.04511 acc 0.98442 val_acc 0.98524

evaluate(model, dataloader)

acc 0.98524: 100%|██████████| 5/5 [00:01<00:00,  4.36it/s]

Guardar nuestros modelos, ya sea el modelo entero o solo su state_dict, es la forma más directa y sencilla de guardar cualquier modelo que hagamos. Sin embargo, tiene ciertas limitaciones. Por un lado, como ya hemos visto, necesitamos la definición del modelo tanto a la hora de entrar como en producción. Esto no solo es engorroso y poco flexible, si no que sólo funcionará en entornos Python con Pytorch instalado. Si bien ésto es suficiente para, por ejemplo, poner nuestros modelos a trabajar en un servidor Flask), en muchas ocasiones necesitaremos ejecutar nuestras redes neuronales en otros entornos (smartphones, aplicaciones web, IoT, ...) en las que usaremos otros lenguajes de programación. Para ello, Pytorch nos permite exportar nuestro modelo en vez de simplemente guradarlo.

Exportando modelos

Torchscript

Torchscript es una representación intermedia de un modelo de Pytorch que puede ejecutarse en diferentes entornos sin la necesidad de Python, por ejemplo en C++. Un modelo de Pytorch exportado en torchscript contiene los pesos de la red así como su definición (todas las operaciones que aplicaremos a un tensor desde la entrada hasta la salida). Tenemos dos maneras de exportar un modelo con torchscript:

• tracing: Dada un entrada, se genera una representación del modelo de manera dinámica registrando todas las operaciones aplicadas al tensor hasta la salida. En este caso no seremos capaces de capturar diferentes caminos en nuestra red (control flow). Es la alternativa más eficiente, pero menos flexible.
• scripting: Genera la representación intermedia de nuestra red neuronal directamente a partir del análisis de la misma, siendo capaz de capturar de manera fiel cualquier ramificación en la misma. No es tan eficiente, pero si más flexible.

# tracing

x = torch.rand(32, 1, 28, 28)
traced_model = torch.jit.trace(model.cpu(), x)
traced_model.save('model.zip')

loaded_model = torch.jit.load('model.zip')
evaluate(loaded_model, dataloader)

acc 0.98524: 100%|██████████| 5/5 [00:01<00:00,  4.38it/s]

# scripting

scripted_model = torch.jit.script(model.cpu())
scripted_model.save('model.zip')

loaded_model = torch.jit.load('model.zip')
evaluate(loaded_model, dataloader)

acc 0.98524: 100%|██████████| 5/5 [00:01<00:00,  4.42it/s]

Exportar nuestro modelo nos aporta varias ventajas:

• Ahora nuestro modelo puede ejecutarse en cualquier entorno capaz de interpretar la representación intermedia generada por torchscript, independientemente del hardware o software utilizado.
• Nuestro modelo contiene los pesos y la definición de las operaciones, evitando tener que guardar código extra.
• Esta representación intermedia puede ser optimizada de manera independiente, haciendo que nuestros modelos sean más rápidos.

La principal desventaja es que al estar "traduciendo" Python a otro lenguaje, es posible que no todas las operaciones que queramos hacer estén soportadas.

⚡ Aprende más sobre Torchscriptaquí.

Un consejo a tener en cuenta a la hora de exportar nuestros modelos que nos puede hacer la vida más fácil cuando pongamos nuestros modelos en producción, es incluir cualquier pre- o post-procesado de datos necesarios en el mismo modelo. Durante el entrenamiento no lo hacemos por motivos de eficiencia (queremos que nuestra GPU entrene la red lo más rápido posible mientras la CPU procesa cada batch de datos de manera paralela), pero estos procesados pueden ser costosos en producción (a veces incluso imposibles de realizar) por lo que incluirlos en el modelo es generalmente una buena idea.

En este caso incluimos la normalización y preparación de los datos en un pre-procesado y calculamos una distribución de probabilidad sobre las salida con un post-procesado, que además también nos devuelve la clase con mayor probabilidad.

class Preprocessing(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        # esperamos un batch de imágenes sin normalizar
        # normalización
        x = (x / 255.)
        x = (x - 0.1307) / 0.3081
        # dimsensiones -> [bs, c, h, w]
        x = x.unsqueeze(1)
        # en imágenes en color, haríamos un `permute`
        return x

class Postprocessing(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.softmax = torch.nn.Softmax(dim=1)
    def forward(self, x) :
        # devolvemos distribución de probabilidad
        # y clase con mayor probabilidad
        return self.softmax(x), torch.argmax(x, dim=1)

final_model = torch.nn.Sequential(
    Preprocessing(),
    model.cpu(),
    Postprocessing()
)

scripted_model = torch.jit.script(final_model)
scripted_model.save('model.zip')

Ahora nuestro modelo acepta lotes de imágenes sin normalizar y con las dimensiones más comunes (alto, ancho) y devuelve una distribución de probabilidad. Es nuestro modelo quien se encarga del procesado.

def script_evaluate(model, dataloader):
    model = torch.jit.load(model)
    model.eval()
    bar = tqdm(dataloader['test'])
    acc = []
    with torch.no_grad():
        for batch in bar:
            X, y = batch
            # desnormalizar
            X = (X*0.3081 + 0.1307)*255
            # quitar dimensión canales
            X = X.squeeze(1)
            # el modelo pre-procesa
            y_hat, label = model(X)
            acc.append((y == label).sum().item() / len(y))
            bar.set_description(f"acc {np.mean(acc):.5f}")

script_evaluate("model.zip", dataloader)

acc 0.98524: 100%|██████████| 5/5 [00:02<00:00,  2.02it/s]

ONNX

Si bien ahora nuestro modelo es capaz de ser importado y ejecutado de forma más flexible y eficiente, seguimos limitados por el hecho de necesitar Pytorch (o la librería Torchscript en C++) para ello. Por este motivo, Pytorch nos ofrece una última manera de exportar nuestros modelos a otra forma de representación intermedia conocida como ONNX. Éste es un formato abierto con el espíritu de convertirse en un estándar de representación de redes neuronales. La gran mayoría de librerías y frameworks de deep learning soportan este formato, lo que implica que, por ejemplo, podemos entrenar un modelo en Pytorch, exportarlo en formato ONNX e importarlo en Tensorflow para ponerlo en producción (aunque ONNX también ofrece soluciones optimizadas para ello: ONNX Runtime). Así pues, exportar nuestros modelos a formato ONNX nos proporcionará la máxima flexibilidad, permitiéndonos, entre muchas otras cosas, ejecutar nuestras redes neuronales en entornos tales como navegadores web o IoT. Por contra, esta librería es la menos flexible en cuanto a cantidad y tipo de operaciones que podemos exportar, lo cual puede imponer unas restricciones muy grandes sobre nuestros modelos (aunque cada vez se soportan más).

x = torch.rand(32, 1, 28, 28)
y = model.cpu()(x)

# exportamos el modelo
torch.onnx.export(model,                     # el modelo
                  x,                         # un ejemplo del input
                  "model.onnx",              # el nombre del archivo para guardar el modelo
                  export_params=True,        # guardar los pesos de la red
                  opset_version=10,          # versión de ONNX
                  do_constant_folding=True,  # optimizaciones
                  input_names = ['input'],   # nombre de los inputs
                  output_names = ['output'], # nombre de los outputs
                  dynamic_axes={'input' : {0 : 'batch_size'},    # ejes con longitud variable (para poder usar diferentes tamaños de batch)
                                'output' : {0 : 'batch_size'}})

Para poder ejecutar nuestro modelo necesitamos la librería de ONNX Runtime para Python, que puedes instalar con el comando pip install onnxruntime.

import onnxruntime

def onnx_evaluate(model, dataloader):
    # cargarmos el modelo
    ort_session = onnxruntime.InferenceSession(model)
    bar = tqdm(dataloader['test'])
    acc = []
    with torch.no_grad():
        for batch in bar:
            X, y = batch
            X, y = X.numpy(), y.numpy()
            # generamos los inputs
            ort_inputs = {ort_session.get_inputs()[0].name: X}
            # extraemos los outputs
            ort_outs = ort_session.run(None, ort_inputs)[0]
            acc.append((y == np.argmax(ort_outs, axis=1)).mean())
            bar.set_description(f"acc {np.mean(acc):.5f}")

onnx_evaluate("model.onnx", dataloader)

acc 0.98524: 100%|██████████| 5/5 [00:01<00:00,  3.50it/s]

De nuevo, te recomiendo incluir tanto pre- como post-procesado como parte del modelo para evitar problemas más adelante.

⚡ Aprende más sobre ONNXaquí.

Resumen

En este post hemos visto las diferentes manera que Pytorch nos ofrece a la hora de guardar y exportar nuestros modelos. Si en nuestro entorno de producción podemos usar Python e instalar Pytorch, entonces podemos simplemente guardar el state_dict de nuestro modelo, o incluso el modelo completo, y luego cargarlo en la aplicación. Ten en cuenta que necesitarás la definición de tu red neuronal (en nuestro ejemplo, la clase Model) que define las capas y operaciones que se aplican dada una entrada. Sin embargo, una opción más eficiente, consiste en exportar nuestro modelo con torch.jit.trace o torch.jit.script ya que luego podremos cargar nuestro modelos sin necesidad de arrastrar código y, además, podremos ejecutarlo en otros entornos como C++. Por último, si no puedes usar Python en tu aplicación, es muy posible que puedas usar ONNX, un formato estándar al que Pytorch nos permite también exportar.