Training CNN on PYNQ z2

I want to use pynq z2 for a CNN task by training the CNN on PC/google colab and then loading the trained(saved) network h5 file for the final classification task but it had some version issues.
So, I switched to training the CNN on pynq itself. But, it seems to ignore the batch size and take the whole train set for 10 epochs.

Training on pynq:

Training on Colab:

This makes the training time so long.
Please help

1 Like

Hi there,

What are the version differences? Your best bet is probably to sort that out rather than training on the Z2 if you want speed, it will be much slower than a dedicated PC/server machine regardless. You could try to upgrade/downgrade the tensorflow version on one of your environments, or look into saving it with onnx, Iā€™m sure tensorflow ought to have hooks for that.

On the epochs issue you might have to add a steps_per_epoch parameter to your fit() function.