[Training] Replace training example model with Conv model gives Segmentation fault

@AdamLouly is already working on this.

Since Feb 10, 2023.

### Describe the issue

Based on the training examples provided in the training_api_demo/mnist_training_example.ipynb, I added conv and max_pool2d into the example model. I tried to train the new model and got segfault.

[root@5CD2160W5Q 20230210]# python3 test.py
2023-02-10 20:27:29.189827112 [I:onnxruntime:Default, reshape_fusion.cc:53 ApplyImpl] Total fused reshape node count: 0
2023-02-10 20:27:29.189872315 [I:onnxruntime:Default, concat_slice_elimination.cc:36 ApplyImpl] Total fused concat node count: 0
2023-02-10 20:27:29.189949359 [I:onnxruntime:Default, reshape_fusion.cc:53 ApplyImpl] Total fused reshape node count: 0
2023-02-10 20:27:29.189956011 [I:onnxruntime:Default, concat_slice_elimination.cc:36 ApplyImpl] Total fused concat node count: 0
2023-02-10 20:27:29.192518618 [W:onnxruntime:Default, checkpoint.cc:187 OrtSaveInternal] Checkpoint directory exists - data may be overwritten.

Segmentation fault

### To reproduce

python version: 3.8.12
Versions of onnxruntime-trainingI've tried and failed, downloaded from https://download.onnxruntime.ai/ :
cu116: 1.15.0.dev20230207001
cpu: 1.15.0.dev20230203001 to 1.15.0.dev20230209001

Here is the modified code based on the example provided from https://github.com/microsoft/onnxruntime-training-examples/blob/master/on_device_training/training_api_demo/mnist_training_example.ipynb

`

import onnxruntime.training.onnxblock as onnxblock
from onnxruntime.training.api import CheckpointState, Module, Optimizer
from onnxruntime import InferenceSession
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
import torch
import onnx
import io
import netron
import evaluate
from torch import nn
import torch.nn.functional as F

# Define a convolutional neural network
class DigitRecognitionNet(nn.Module):
def __init__(self):
super(DigitRecognitionNet, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.fc1 = nn.Linear(320, 50)
# self.fc1 = nn.Linear(784, 50)
self.fc2 = nn.Linear(50, 10)

def forward(self, input1):
# x = input1
x = F.relu(F.max_pool2d(self.conv1(input1), 2))
x = F.relu(F.max_pool2d(self.conv2(x), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = self.fc2(x)
x = F.log_softmax(x, dim=1)
return x

# Create a MNISTNet instance.
# Generate a random input.
device = "cpu"

pt_model = DigitRecognitionNet()
model_inputs = (torch.randn((64, 1, 28, 28), device=device), )
# model_inputs = (torch.randn((64, 784), device=device), )

model_outputs = pt_model(*model_inputs)
if isinstance(model_outputs, torch.Tensor):
model_outputs = [model_outputs]

dynamic_axes = {}
input_names = []
output_names = []

for i, model_input in enumerate(model_inputs):
input_name = f"input-{i}"
input_names.append(input_name)
dynamic_axes[input_name] = {}
for dim_idx in range(len(model_input.shape)):
dynamic_axes[input_name].update(
{dim_idx: f"{input_name}_dim{dim_idx}"})

for i, model_output in enumerate(model_outputs):
output_name = f"output-{i}"
output_names.append(output_name)
dynamic_axes[output_name] = {}
for dim_idx in range(len(model_output.shape)):
dynamic_axes[output_name].update(
{dim_idx: f"{output_name}_dim{dim_idx}"})

f = io.BytesIO()
torch.onnx.export(
pt_model,
model_inputs,
f,
input_names=input_names,
output_names=output_names,
opset_version=14,
do_constant_folding=False,
training=torch.onnx.TrainingMode.TRAINING,
dynamic_axes=dynamic_axes,
export_params=True,
keep_initializers_as_inputs=False,
)
onnx_model = onnx.load_model_from_string(f.getvalue())

# Creating a class with a Loss function.
class ModelWithLoss(onnxblock.TrainingModel):
def __init__(self):
super(ModelWithLoss, self).__init__()
self.loss = onnxblock.loss.CrossEntropyLoss()

def build(self, output_name):
return self.loss(output_name), output_name

# Build the onnx model with loss
simple_model = ModelWithLoss()

# Building training graph and eval graph.
with onnxblock.onnx_model(onnx_model) as accessor:
_ = simple_model(onnx_model.graph.output[0].name)
eval_model = accessor.eval_model

# Building the optimizer graph
optimizer = onnxblock.optim.AdamW()
with onnxblock.onnx_model() as accessor:
_ = optimizer(simple_model.parameters())
optimizer_model = accessor.model

# Saving all the files to use them later for the training.
trainable_params, non_trainable_params = simple_model.parameters()
onnxblock.save_checkpoint((trainable_params, non_trainable_params),
"data/checkpoint.ckpt")
onnx.save(onnx_model, "data/training_model.onnx")
onnx.save(optimizer_model, "data/optimizer.onnx")
onnx.save(eval_model, "data/eval_model.onnx")

batch_size = 64
train_kwargs = {'batch_size': batch_size}
test_batch_size = 1000
test_kwargs = {'batch_size': test_batch_size}

transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))])

dataset1 = datasets.MNIST("/data",
train=True,
download=True,
transform=transform)
dataset2 = datasets.MNIST("/data",
train=False,
download=True,
transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)

# Create checkpoint state.
state = CheckpointState("data/checkpoint.ckpt")

# Create module.
model = Module("data/training_model.onnx", state, "data/eval_model.onnx")

# Create optimizer.
optimizer = Optimizer("data/optimizer.onnx", model)

# Util function to convert logits to predictions.
def get_pred(logits):
return np.argmax(logits, axis=1)

# Training Loop :
def train(epoch):
model.train()
losses = []
for batch_idx, (data, target) in enumerate(train_loader):
forward_inputs = [data.numpy(), target.numpy().astype(np.int32)]
train_loss, _ = model(forward_inputs)
optimizer.step()
model.lazy_reset_grad()
losses.append(train_loss)

print(f'Epoch: {epoch+1},Train Loss: {sum(losses)/len(losses):.4f}')

# Test Loop :
def test(epoch):
model.eval()
losses = []
metric = evaluate.load('accuracy')

for batch_idx, (data, target) in enumerate(train_loader):
forward_inputs = [
data.numpy(), #.reshape(len(data), 784).numpy(),
target.numpy().astype(np.int32)
]
test_loss, logits = model(forward_inputs)
metric.add_batch(references=target, predictions=get_pred(logits))
losses.append(test_loss)

metrics = metric.compute()
print(
f'Epoch: {epoch+1}, Test Loss: {sum(losses)/len(losses):.4f}, Accuracy : {metrics["accuracy"]:.2f}'
)

for epoch in range(5):
train(epoch)
test(epoch)

model.export_model_for_inferencing("data/inference_model.onnx", ["output-0"])
session = InferenceSession('data/inference_model.onnx',
providers=['CPUExecutionProvider'])

# getting one example from test list to try inference.
data = next(iter(test_loader))[0][0]

input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
output = session.run([output_name], {input_name: data.reshape(1, 784).numpy()})

# plotting the picture
plt.imshow(data[0], cmap='gray')
plt.savefig('digit.png')

print("Predicted Label : ", get_pred(output[0]))
`

### Urgency

_No response_

### ONNX Runtime Installation

Other / Unknown

### ONNX Runtime Version or Commit ID

1.13.1

### PyTorch Version

1.11.0

### Execution Provider

Default CPU

### Execution Provider Library Version

_No response_