This week was less about improving the model and more about making the whole system work end to end.
The main goal was simple: take a processed embryo, pass it through ESTGEL, compute a loss, backpropagate, and save a checkpoint without anything breaking.
One important part was building the training wrapper around the dataset. The processed EPIC data only contains embryos, so I added a small label wrapper that attaches the correct class when each sample is loaded. This keeps the preprocessing pipeline separate from the training logic.
I also added a stratified train-validation split so the classes stay balanced, and class-weighted cross-entropy to handle imbalance automatically. On the model side, I made most training settings configurable through command-line arguments, which makes it much easier to run different experiments without changing the code every time.
The training loop itself is fairly straightforward: load an embryo, run a forward pass, compute loss, backpropagate, clip gradients, and update the weights. I also added checkpointing for both the latest and best model states, along with the config and training history.
The first smoke test was small, just four embryos for one epoch, but it was enough to confirm that the full pipeline works. The accuracy was around 50%, which is expected for a random initialization on such a tiny run, but the important part is that training, validation, and checkpoint saving all completed successfully.