[Python] I really need help in understanding what is the core problem of this error, I've...

I am currently building a ResNet-152 architecture to classify diatom images, using MLX , but I am quite new to it.

I will leave the code for the BottleNeck and the ResNet architecture plus the training code with the error, if somebody can answer this and explain I would really appreciate it!

import mlx.nn as nn


import mlx.core as mx


class BottleNeck(nn.Module):expansion = 4 # class-level constant, NOT self.expansion

def __init__(self, in_channels, out_channels, identity_downsample=None, stride=1):
super().__init__()

self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
self.bn1 = nn.BatchNorm(out_channels)

self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1)
self.bn2 = nn.BatchNorm(out_channels)

self.conv3 = nn.Conv2d(out_channels, out_channels * BottleNeck.expansion, kernel_size=1, stride=1, padding=0)
self.bn3 = nn.BatchNorm(out_channels * BottleNeck.expansion)

self.relu = nn.ReLU()
self.identity_downsample = identity_downsample

def __call__(self, x):
identity = x

x = self.relu(self.bn1(self.conv1(x)))
x = self.relu(self.bn2(self.conv2(x)))
x = self.bn3(self.conv3(x))

if self.identity_downsample is not None:
identity = self.identity_downsample(identity)

x += identity
return self.relu(x)`

import mlx.nn as nn
import mlx.core as mx

class ResNet(nn.Module):
def __init__(self, BottleNeck, layers, image_channels, num_classes):
super().__init__()
self._in_channels = 64 # use underscore to avoid optimizer confusion

self.conv1 = nn.Conv2d(image_channels, self._in_channels, kernel_size=7, stride=2, padding=3)
self.bn1 = nn.BatchNorm(self._in_channels)
self.relu = nn.ReLU()
self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

self.layer1 = self.make_layers(BottleNeck, layers[0], out_channels=64, stride=1)
self.layer2 = self.make_layers(BottleNeck, layers[1], out_channels=128, stride=2)
self.layer3 = self.make_layers(BottleNeck, layers[2], out_channels=256, stride=2)
self.layer4 = self.make_layers(BottleNeck, layers[3], out_channels=512, stride=2)

self.fcl = nn.Linear(512 * 4, num_classes)

def make_layers(self, BottleNeck, num_residual_blocks, out_channels, stride):
identity_downsample = None
if stride != 1 or self._in_channels != out_channels * 4:
identity_downsample = nn.Sequential(
nn.Conv2d(self._in_channels, out_channels * 4, kernel_size=1, stride=stride),
nn.BatchNorm(out_channels * 4)
)

layers = [BottleNeck(self._in_channels, out_channels, identity_downsample, stride)]
self._in_channels = out_channels * 4

for _ in range(1, num_residual_blocks):
layers.append(BottleNeck(self._in_channels, out_channels))

return nn.Sequential(*layers)

def __call__(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.pool1(x)

x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)

x = mx.mean(x, axis=(1, 2)) # global average pooling
x = self.fcl(x)
return x


from tqdm import tqdm
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim

num_classes = len(set(label_dict.values()))
model = ResNet(BottleNeck, [3, 8, 36, 3], image_channels=3, num_classes=num_classes)

optimizer = optim.Adam(1e-4)
loss_function = nn.losses.cross_entropy

# Run one batch to initialize model parameters
example_input, _ = next(iter(train_loader))
example_input = mx.array(example_input.permute(0, 2, 3, 1).numpy())
_ = model(example_input)

# Initialize optimizer with model parameters
optimizer.init(model)

losses = []

for epoch in range(50):
print(f"Training epoch {epoch + 1}")
running_loss = 0.0

for inputs, labels in tqdm(train_loader):
inputs = mx.array(inputs.permute(0, 2, 3, 1).numpy()) # Convert to NHWC
labels = mx.array(labels.numpy())

def forward(x, y):
outputs = model(x)
return loss_function(outputs, y).mean()

loss, grads = mx.value_and_grad(forward)(inputs, labels)
optimizer.update(model, grads)

running_loss += loss.item()

epoch_loss = running_loss / len(train_loader)
losses.append(epoch_loss)
print(f"Loss: {epoch_loss:.4f}")


MLX is trying to initialize model parameters as integers instead of tensors but I can't find the problem in my code.

Continue reading...