AutoEncoders

Collection of Autoencoder models

Fashion MNIST

cfg = OmegaConf.load('../config/data/image/image.yaml')
dm = instantiate(cfg, name='fashion_mnist', data_dir='../data/image/')
dm.prepare_data()
dm.setup()
print(dm.num_classes)

[22:06:36] INFO - Init ImageDataModule for fashion_mnist
[22:06:56] INFO - split train into train/val [0.8, 0.2]
[22:06:56] INFO - train: 48000 val: 12000, test: 10000

10

dm.show_grid(5,5)

print(dm.label_names)

['T - shirt / top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

ConvNet

cfg = OmegaConf.load('../config/model/image/convnetx_adam.yaml')
# nnet = instantiate(cfg.nnet, num_classes=dm.num_classes)
# optimizer = instantiate(cfg.optimizer)
# scheduler = instantiate(cfg.scheduler)

# model = ConvNetX(nnet, dm.num_classes, optimizer, scheduler)
model = instantiate(cfg, num_classes=dm.num_classes)

[22:06:57] INFO - ConvNetX: init
[22:06:57] INFO - Classifier: init
/Users/slegroux/miniforge3/envs/nimrod/lib/python3.11/site-packages/lightning/pytorch/utilities/parsing.py:208: Attribute 'nnet' is an instance of `nn.Module` and is already saved during checkpointing. It is recommended to ignore them using `self.save_hyperparameters(ignore=['nnet'])`.

MAX_EPOCHS = 5
dm.batch_size = 256
print(dm.batch_size)
# lr = 0.4

trainer = Trainer(
    max_epochs=MAX_EPOCHS,
    logger=CSVLogger("logs", name="fashion_mnist_convnet"),
    callbacks = [LearningRateMonitor(logging_interval="step")],
    check_val_every_n_epoch=1,
    log_every_n_steps=1
    )

GPU available: True (mps), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs

256

LR Finder

tuner = Tuner(trainer)
lr_finder = tuner.lr_find(
    model,
    datamodule=dm,
    min_lr=1e-6,
    max_lr=1.0,
    num_training=100,  # number of iterations
    # attr_name="optimizer.lr",
)
fig = lr_finder.plot(suggest=True)
plt.show()
print(f"Suggested learning rate: {lr_finder.suggestion()}")

[22:09:59] INFO - Optimizer: Adam (
Parameter Group 0
    amsgrad: False
    betas: (0.9, 0.999)
    capturable: False
    differentiable: False
    eps: 1e-08
    foreach: None
    fused: None
    lr: 0.06
    maximize: False
    weight_decay: 0
)
[22:09:59] INFO - Scheduler: <torch.optim.lr_scheduler.ReduceLROnPlateau object>
/Users/slegroux/miniforge3/envs/nimrod/lib/python3.11/site-packages/lightning/pytorch/core/optimizer.py:316: The lr scheduler dict contains the key(s) ['monitor', 'strict'], but the keys will be ignored. You need to call `lr_scheduler.step()` manually in manual optimization.
/Users/slegroux/miniforge3/envs/nimrod/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/Users/slegroux/miniforge3/envs/nimrod/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

`Trainer.fit` stopped: `max_steps=100` reached.
Learning rate set to 9.120108393559098e-06
Restoring states from the checkpoint path at /Users/slegroux/Projects/nimrod/nbs/.lr_find_5828b967-b82f-4e1a-bda1-997d275f4d03.ckpt
Restored all states from the checkpoint at /Users/slegroux/Projects/nimrod/nbs/.lr_find_5828b967-b82f-4e1a-bda1-997d275f4d03.ckpt

Suggested learning rate: 9.120108393559098e-06

# print(f"lr: {model.lr}, bs: {dm.batch_size}")

lr: 9.120108393559098e-06, bs: 256

Fit

cfg.optimizer.lr = 0.4
print(OmegaConf.to_yaml(cfg))

model = instantiate(cfg)
trainer.fit(model, dm.train_dataloader(), dm.val_dataloader())

[22:10:36] INFO - ConvNetX: init
[22:10:36] INFO - Classifier: init

_target_: nimrod.models.conv.ConvNetX
num_classes: 10
nnet:
  _target_: nimrod.models.conv.ConvNet
  n_features:
  - 1
  - 8
  - 16
  - 32
  - 64
  num_classes: ${..num_classes}
  kernel_size: 3
  bias: null
  normalization:
    _target_: hydra.utils.get_class
    path: torch.nn.BatchNorm2d
  activation:
    _target_: hydra.utils.get_class
    path: torch.nn.ReLU
optimizer:
  _target_: torch.optim.Adam
  _partial_: true
  lr: 0.4
scheduler:
  _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
  _partial_: true
  mode: min
  factor: 0.1
  patience: 5

/Users/slegroux/miniforge3/envs/nimrod/lib/python3.11/site-packages/lightning/pytorch/utilities/parsing.py:208: Attribute 'nnet' is an instance of `nn.Module` and is already saved during checkpointing. It is recommended to ignore them using `self.save_hyperparameters(ignore=['nnet'])`.
[22:10:54] INFO - Optimizer: Adam (
Parameter Group 0
    amsgrad: False
    betas: (0.9, 0.999)
    capturable: False
    differentiable: False
    eps: 1e-08
    foreach: None
    fused: None
    lr: 0.4
    maximize: False
    weight_decay: 0
)
[22:10:54] INFO - Scheduler: <torch.optim.lr_scheduler.ReduceLROnPlateau object>

  | Name         | Type               | Params | Mode 
------------------------------------------------------------
0 | loss         | CrossEntropyLoss   | 0      | train
1 | train_acc    | MulticlassAccuracy | 0      | train
2 | val_acc      | MulticlassAccuracy | 0      | train
3 | test_acc     | MulticlassAccuracy | 0      | train
4 | train_loss   | MeanMetric         | 0      | train
5 | val_loss     | MeanMetric         | 0      | train
6 | test_loss    | MeanMetric         | 0      | train
7 | val_acc_best | MaxMetric          | 0      | train
8 | nnet         | ConvNet            | 30.3 K | train
------------------------------------------------------------
30.3 K    Trainable params
0         Non-trainable params
30.3 K    Total params
0.121     Total estimated model params size (MB)
34        Modules in train mode
0         Modules in eval mode

[22:11:24] INFO - scheduler is an instance of Reduce plateau

[22:11:54] INFO - scheduler is an instance of Reduce plateau

[22:12:24] INFO - scheduler is an instance of Reduce plateau

[22:12:55] INFO - scheduler is an instance of Reduce plateau

[22:13:25] INFO - scheduler is an instance of Reduce plateau
`Trainer.fit` stopped: `max_epochs=5` reached.

########################
csv_path = f"{trainer.logger.log_dir}/metrics.csv"
metrics = pd.read_csv(csv_path)
metrics.head()

#########################
plt.figure()
plt.plot(metrics['step'], metrics['train/loss_step'], 'b.-')
plt.plot(metrics['step'], metrics['val/loss'],'r.-')
plt.figure()
plt.plot(metrics['step'], metrics['lr-Adam'], 'g.-')
plt.show()

trainer.test(model, dm.test_dataloader())

/Users/slegroux/miniforge3/envs/nimrod/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃        Test metric        ┃       DataLoader 0        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│         test/acc          │     0.803600013256073     │
│         test/loss         │    0.5951264500617981     │
└───────────────────────────┴───────────────────────────┘

[{'test/loss': 0.5951264500617981, 'test/acc': 0.803600013256073}]

FC Autoencoder

---

AutoEncoder

 AutoEncoder (encoder:torch.nn.modules.module.Module,
              decoder:torch.nn.modules.module.Module)

A modular autoencoder with configurable encoder and decoder

	Type	Details
encoder	Module	Encoder layer
decoder	Module	Decoder layer

class LinearEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))

    def forward(self, x):
        return self.l1(x)


class LinearDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))

    def forward(self, x):
        return self.l1(x)

enc = LinearEncoder()
dec = LinearDecoder()
a = AutoEncoder(enc, dec)
batch = torch.rand((5, 3, 28*28))
y = a(batch)
print(y.shape)

torch.Size([5, 3, 784])

ds = ImageDataset(name='fashion_mnist', data_dir='../data/image/')
dl = DataLoader(ds, batch_size=3)
b = next(iter(dl))
print(f" X: {b[0].shape}, Y: {b[1].shape}")

 X: torch.Size([3, 1, 28, 28]), Y: torch.Size([3])

acfg = OmegaConf.load('../config/data/image/image.yaml')
dm = instantiate(cfg, name='fashion_mnist', data_dir='../data/image/')
dm.prepare_data()
dm.setup()
# print(f"num classes: {dm.num_classes}, bs: {dm.batch_size}, labels: {dm.label_names}" if dm.label_names else f"num classes: {dm.num_classes}")

[14:58:07] INFO - Init ImageDataModule for fashion_mnist
[14:58:22] INFO - split train into train/val [0.8, 0.2]
[14:58:22] INFO - train: 48000 val: 12000, test: 10000

device = get_device()
print(f"Device: {device}")
enc = LinearEncoder()
dec =LinearDecoder()
model = AutoEncoder(enc, dec).to(device)

[14:55:26] INFO - Using device: mps

Device: mps

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)


N_EPOCHS = 5

for epoch in range(N_EPOCHS):
    i = 0
    model.train()
    for images, labels in dm.train_dataloader():
        optimizer.zero_grad()
        images, labels = images.to(device), labels.to(device)
        # B x C x H x W -> B x C x L
        images = images.view(-1, images.size(2) * images.size(3))
        outputs = model(images)
        # output should be as close to input as possible
        loss = criterion(outputs, images)        
        loss.backward()
        optimizer.step()

    model.eval()
    with torch.no_grad():
        total_loss, epoch_step = 0, 0
        for images, labels in dm.val_dataloader():
            images, labels = images.to(device), labels.to(device)
            images = images.view(-1, images.size(2) * images.size(3))
            outputs = model(images)
            eval_loss = criterion(outputs, images)
            epoch_len = len(images)
            epoch_step += epoch_len
            total_loss += eval_loss.item() * epoch_len
    logger.info(f"Epoch: {epoch}, len: {epoch_len}, Loss: {total_loss / epoch_step:.3f}")

[14:56:24] INFO - Epoch: 0, len: 32, Loss: 0.025
[14:56:27] INFO - Epoch: 1, len: 32, Loss: 0.025
[14:56:31] INFO - Epoch: 2, len: 32, Loss: 0.025
[14:56:35] INFO - Epoch: 3, len: 32, Loss: 0.025
[14:56:39] INFO - Epoch: 4, len: 32, Loss: 0.025

x, y = next(iter(dm.train_dataloader()))
print(f" X: {x.shape}, Y: {y.shape}")
x = x.to(device)
B, C, H, W = x.shape
x_hat = model(x.view(-1, H * W)).view(-1, C, H, W)
print(f" X_hat: {x_hat.shape}")
idx = 0
n_rows, n_cols = 1, 2
fig, axs = plt.subplots(n_rows, n_cols, figsize=(10, 10))
axs[0].imshow(x[idx].permute(1, 2, 0).cpu().numpy(), cmap='gray')
axs[1].imshow(x_hat[idx].permute(1, 2, 0).detach().cpu().numpy(), cmap='gray')
plt.show()

 X: torch.Size([64, 1, 28, 28]), Y: torch.Size([64])
 X_hat: torch.Size([64, 1, 28, 28])

ConvNet Autoencoder

class ConvEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        layers = nn.ModuleList()
        # X -> B,C,28,28
        layers.append(nn.ZeroPad2d(2)) # X -> B,C,32,32
        layers.append(ConvLayer(1,2, normalization=None)) # 16 x 16
        layers.append(ConvLayer(2,4, normalization=None)) # 8 x 8
        self._nnet = nn.Sequential(*layers)

    def forward(self, x:torch.Tensor)->torch.Tensor:
        return self._nnet(x)

class ConvDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        layers = nn.ModuleList()
        layers.append(DeconvLayer(4,2, normalization=None)) # 16 x 16
        layers.append(DeconvLayer(2,1, normalization=None, activation=None)) # 32 x 32
        layers.append(nn.ZeroPad2d(-2)) # 28 x 28
        layers.append(nn.Sigmoid())
        self._nnet = nn.Sequential(*layers)

    def forward(self, x:torch.Tensor)->torch.Tensor:
        return self._nnet(x)

device = get_device()
print(f"Device: {device}")
enc = ConvEncoder()
dec = ConvDecoder()
model = AutoEncoder(enc, dec).to(device)

[15:32:25] INFO - Using device: mps

Device: mps

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)


N_EPOCHS = 5

for epoch in range(N_EPOCHS):
    i = 0
    model.train()
    for images, labels in dm.train_dataloader():
        optimizer.zero_grad()
        images, labels = images.to(device), labels.to(device)

        # images = images.view(-1, images.size(2) * images.size(3))
        outputs = model(images)
        # output should be as close to input as possible
        loss = criterion(outputs, images)        
        loss.backward()
        optimizer.step()

    model.eval()
    with torch.no_grad():
        total_loss, epoch_step = 0, 0
        for images, labels in dm.val_dataloader():
            images, labels = images.to(device), labels.to(device)
            # images = images.view(-1, images.size(2) * images.size(3))
            outputs = model(images)
            eval_loss = criterion(outputs, images)
            epoch_len = len(images)
            epoch_step += epoch_len
            total_loss += eval_loss.item() * epoch_len
    logger.info(f"Epoch: {epoch}, len: {epoch_len}, Loss: {total_loss / epoch_step:.3f}")

[15:34:38] INFO - Epoch: 0, len: 32, Loss: 0.012
[15:34:43] INFO - Epoch: 1, len: 32, Loss: 0.011
[15:34:48] INFO - Epoch: 2, len: 32, Loss: 0.011
[15:34:53] INFO - Epoch: 3, len: 32, Loss: 0.011
[15:34:59] INFO - Epoch: 4, len: 32, Loss: 0.011

x, y = next(iter(dm.train_dataloader()))
print(f" X: {x.shape}, Y: {y.shape}")
x = x.to(device)
B, C, H, W = x.shape
x_hat = model(x)
print(f" X_hat: {x_hat.shape}")
idx = 0
n_rows, n_cols = 1, 2
fig, axs = plt.subplots(n_rows, n_cols, figsize=(10, 10))
axs[0].imshow(x[idx].permute(1, 2, 0).cpu().numpy(), cmap='gray')
axs[1].imshow(x_hat[idx].permute(1, 2, 0).detach().cpu().numpy(), cmap='gray')
plt.show()

 X: torch.Size([64, 1, 28, 28]), Y: torch.Size([64])
 X_hat: torch.Size([64, 1, 28, 28])

AutoEncoder_X

#show_doc(AutoEncoderPL.forward)

---

AutoEncoderPL.forward

 AutoEncoderPL.forward (x:torch.Tensor)

Forward pass of the AutoEncoder model.

	Type	Details
x	Tensor	Tensor B x L
Returns	Tensor	Reconstructed input tensor of shape B x L

autoencoder_pl = AutoEncoderPL(a)
b = torch.rand((5,28*28))
y = autoencoder_pl(b)
print(y.shape)

torch.Size([5, 784])