Change the TrainableModel::summary API to return `ModelSummary`

Kotlin / kotlindl

High-level Deep Learning Framework written in Kotlin and inspired by Keras

Apache License 2.0

1.44k stars 102 forks source link

Change the TrainableModel::summary API to return `ModelSummary` #135

Closed avan1235 closed 2 years ago

avan1235 commented 3 years ago

In current implementation the TrainableModel::summary return the list with the descriptions for the following layers.

This information is readable but doesn't allow for easy transformation of model summary gathered data as it is give as raw String. Additionally, the implementation classes log the provided layers description in more readable way, but it is only printed on stdio and not available at least as pretty formatted string that could bu printed later. The problem is more annoying in jupyter notebooks where we don't have an access to stdio and so the logged pretty summary is not visible.

I would propose creating some simple ModelSummary data class that would be responsible for holding the model summary data and would have some good implementation of toString() method that could be easily printed. It would allow for creating custom printing methods for ModelSummary as well as good keeping of model summary data as a single object instead of List<String>

knok16 commented 2 years ago

What do you think about the next changes?

Classes to represent a model summary

data class ModelSummary(
    val type: String,
    val name: String,
    val layersSummaries: List<LayerSummary>,
    val trainableParamsCount: Long,
    val frozenParamsCount: Long,
    val totalParamsCount: Long
)

data class LayerSummary(
    val name: String,
    val type: String,
    val shape: Shape,
    val paramsCount: Long,
    val inputLayers: List<String>
)

Make TrainableModel::summary return ModelSummary
Remove all the logging from TrainableModel::summary so it will not have any side effects
Provide some functions that will print ModelSummary to console/logger