---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[5], line 137
    135 i = 0
    136 for epoch in range(EPOCHS):
--> 137     for batch in tqdm(dataloaders["train"], desc=f"Epoch {epoch + 1}/{EPOCHS}"):
    138         i += 1
    140         loss = compute_loss(batch)

File ~/.local/lib/python3.10/site-packages/tqdm/std.py:1181, in tqdm.__iter__(self)
   1178 time = self._time
   1180 try:
-> 1181     for obj in iterable:
   1182         yield obj
   1183         # Update and possibly print the progressbar.
   1184         # Note: does not call self.update(1) for speed optimisation.

File ~/.local/lib/python3.10/site-packages/torch/utils/data/dataloader.py:631, in _BaseDataLoaderIter.__next__(self)
    628 if self._sampler_iter is None:
    629     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    630     self._reset()  # type: ignore[call-arg]
--> 631 data = self._next_data()
    632 self._num_yielded += 1
    633 if self._dataset_kind == _DatasetKind.Iterable and \
    634         self._IterableDataset_len_called is not None and \
    635         self._num_yielded > self._IterableDataset_len_called:

File ~/.local/lib/python3.10/site-packages/torch/utils/data/dataloader.py:675, in _SingleProcessDataLoaderIter._next_data(self)
    673 def _next_data(self):
    674     index = self._next_index()  # may raise StopIteration
--> 675     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    676     if self._pin_memory:
    677         data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)

File ~/.local/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:54, in _MapDatasetFetcher.fetch(self, possibly_batched_index)
     52 else:
     53     data = self.dataset[possibly_batched_index]
---> 54 return self.collate_fn(data)

Cell In[5], line 14
     12 def collate_fn(batch):
     13     images = [sample['image'] for sample in batch]
---> 14     images = torch.stack(moondream.vision_encoder.preprocess(images))
     15     images = rearrange(images,
     16                        "b c (h p1) (w p2) -> b (h w) (c p1 p2)",
     17                        p1=14, p2=14)
     19     labels_acc = []

File ~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/.local/lib/python3.10/site-packages/torchvision/transforms/v2/_container.py:51, in Compose.forward(self, *inputs)
     49 needs_unpacking = len(inputs) > 1
     50 for transform in self.transforms:
---> 51     outputs = transform(*inputs)
     52     inputs = outputs if needs_unpacking else (outputs,)
     53 return outputs

File ~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/.local/lib/python3.10/site-packages/torchvision/transforms/v2/_transform.py:50, in Transform.forward(self, *inputs)
     45 needs_transform_list = self._needs_transform_list(flat_inputs)
     46 params = self._get_params(
     47     [inpt for (inpt, needs_transform) in zip(flat_inputs, needs_transform_list) if needs_transform]
     48 )
---> 50 flat_outputs = [
     51     self._transform(inpt, params) if needs_transform else inpt
     52     for (inpt, needs_transform) in zip(flat_inputs, needs_transform_list)
     53 ]
     55 return tree_unflatten(flat_outputs, spec)

File ~/.local/lib/python3.10/site-packages/torchvision/transforms/v2/_transform.py:51, in <listcomp>(.0)
     45 needs_transform_list = self._needs_transform_list(flat_inputs)
     46 params = self._get_params(
     47     [inpt for (inpt, needs_transform) in zip(flat_inputs, needs_transform_list) if needs_transform]
     48 )
     50 flat_outputs = [
---> 51     self._transform(inpt, params) if needs_transform else inpt
     52     for (inpt, needs_transform) in zip(flat_inputs, needs_transform_list)
     53 ]
     55 return tree_unflatten(flat_outputs, spec)

File ~/.local/lib/python3.10/site-packages/torchvision/transforms/v2/_misc.py:165, in Normalize._transform(self, inpt, params)
    164 def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
--> 165     return self._call_kernel(F.normalize, inpt, mean=self.mean, std=self.std, inplace=self.inplace)

File ~/.local/lib/python3.10/site-packages/torchvision/transforms/v2/_transform.py:35, in Transform._call_kernel(self, functional, inpt, *args, **kwargs)
     33 def _call_kernel(self, functional: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any:
     34     kernel = _get_kernel(functional, type(inpt), allow_passthrough=True)
---> 35     return kernel(inpt, *args, **kwargs)

File ~/.local/lib/python3.10/site-packages/torchvision/transforms/v2/functional/_utils.py:31, in _kernel_tv_tensor_wrapper.<locals>.wrapper(inpt, *args, **kwargs)
     20 @functools.wraps(kernel)
     21 def wrapper(inpt, *args, **kwargs):
     22     # If you're wondering whether we could / should get rid of this wrapper,
   (...)
     29     # lost after the first operation due to our own __torch_function__
     30     # logic.
---> 31     output = kernel(inpt.as_subclass(torch.Tensor), *args, **kwargs)
     32     return tv_tensors.wrap(output, like=inpt)

File ~/.local/lib/python3.10/site-packages/torchvision/transforms/v2/functional/_misc.py:65, in normalize_image(image, mean, std, inplace)
     63     image = image.sub_(mean)
     64 else:
---> 65     image = image.sub(mean)
     67 return image.div_(std)

RuntimeError: The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 0


from torch.utils.data import DataLoader
from bitsandbytes.optim import Adam8bit
import math
from einops import rearrange
from tqdm import tqdm

ANSWER_EOS = "
"

# Number of tokens used to represent each image.
IMG_TOKENS = 729

def collate_fn(batch):
    images = [sample['image'] for sample in batch]
    images = torch.stack(moondream.vision_encoder.preprocess(images))
    images = rearrange(images,
                       "b c (h p1) (w p2) -> b (h w) (c p1 p2)",
                       p1=14, p2=14)

    labels_acc = []
    tokens_acc = []

    for sample in batch:
        toks = [tokenizer.bos_token_id]
        labs = [-100] * (IMG_TOKENS + 1)

        for qa in sample['qa']:
            q_t = tokenizer(
                f"\n\nQuestion: {qa['question']}\n\nAnswer:",
                add_special_tokens=False
            ).input_ids
            toks.extend(q_t)
            labs.extend([-100] * len(q_t))

            a_t = tokenizer(
                f" {qa['answer']}{ANSWER_EOS}",
                add_special_tokens=False
            ).input_ids
            toks.extend(a_t)
            labs.extend(a_t)

        tokens_acc.append(toks)
        labels_acc.append(labs)

    max_len = -1
    for labels in labels_acc:
        max_len = max(max_len, len(labels))

    attn_mask_acc = []

    for i in range(len(batch)):
        len_i = len(labels_acc[i])
        pad_i = max_len - len_i

        labels_acc[i].extend([-100] * pad_i)
        tokens_acc[i].extend([tokenizer.eos_token_id] * pad_i)
        attn_mask_acc.append([1] * len_i + [0] * pad_i)

    return (
        images.to(dtype=DTYPE),
        torch.stack([torch.tensor(t, dtype=torch.long) for t in tokens_acc]),
        torch.stack([torch.tensor(l, dtype=torch.long) for l in labels_acc]),
        torch.stack([torch.tensor(a, dtype=torch.bool) for a in attn_mask_acc]),
    )

def compute_loss(batch):
    print(batch)
    images, tokens, labels, attn_mask = batch

    images = images.to(DEVICE)
    tokens = tokens.to(DEVICE)
    labels = labels.to(DEVICE)
    attn_mask = attn_mask.to(DEVICE)

    with torch.no_grad():
        img_embs = moondream.vision_encoder.encoder(images)
        img_embs = moondream.vision_encoder.projection(img_embs)

    tok_embs = moondream.text_model.get_input_embeddings()(tokens)
    inputs_embeds = torch.cat((tok_embs[:, 0:1, :], img_embs, tok_embs[:, 1:, :]), dim=1)

    outputs = moondream.text_model(
        inputs_embeds=inputs_embeds,
        labels=labels,
        attention_mask=attn_mask,
    )

    return outputs.loss

def lr_schedule(step, max_steps):
    x = step / max_steps
    if x < 0.1:
        return 0.1 * LR + 0.9 * LR * x / 0.1
    else:
        return 0.1 * LR + 0.9 * LR * (1 + math.cos(math.pi * (x - 0.1))) / 2

dataloaders = {
    "train": DataLoader(
        datasets["train"],
        batch_size=BATCH_SIZE,
        shuffle=True,
        collate_fn=collate_fn,
    ),
    "val": DataLoader(
        datasets["val"],
        batch_size=BATCH_SIZE,
        collate_fn=collate_fn,
    ),
}

moondream.text_model.train()
moondream.text_model.transformer.gradient_checkpointing_enable()

total_steps = EPOCHS * len(dataloaders["train"]) // GRAD_ACCUM_STEPS
optimizer = Adam8bit(
    [
        {"params": moondream.text_model.parameters()},
    ],
    lr=LR * 0.1,
    betas=(0.9, 0.95),
    eps=1e-6
)

if USE_WANDB:
    import wandb
    wandb.init(
        project="moondream-ft",
        config={
            "EPOCHS": EPOCHS,
            "BATCH_SIZE": BATCH_SIZE,
            "GRAD_ACCUM_STEPS": GRAD_ACCUM_STEPS,
            "LR": LR,
        }
    )

i = 0
for epoch in range(EPOCHS):
    for batch in tqdm(dataloaders["train"], desc=f"Epoch {epoch + 1}/{EPOCHS}"):
        i += 1

        loss = compute_loss(batch)
        loss.backward()

        if i % GRAD_ACCUM_STEPS == 0:
            optimizer.step()
            optimizer.zero_grad()

            lr = lr_schedule(i / GRAD_ACCUM_STEPS, total_steps)
            for param_group in optimizer.param_groups:
                param_group['lr'] = lr

        if i % 8 == 0:
            torch.cuda.empty_cache()

        if i % 100 == 0 and USE_WANDB:
            # Calculate validation loss
            val_loss = 0
            for val_batch in tqdm(dataloaders["val"], desc="Validation"):
                with torch.no_grad():
                    val_loss += compute_loss(val_batch).item()
            val_loss /= len(dataloaders["val"])

        if USE_WANDB:
            wandb.log({
                "loss/train": loss.item(),
                "lr": optimizer.param_groups[0]['lr']
            } | ({"loss/val": val_loss} if i % 100 == 0 else {}))

if USE_WANDB:
    wandb.finish()