This paper presents a measurement of the production cross-section of a $Z$ boson in association with $b$-jets, in proton-proton collisions at $\sqrt{s} = 13$ TeV with the ATLAS experiment at the Large Hadron Collider using data corresponding to an integrated luminosity of 35.6 fb$^{-1}$. Inclusive and differential cross-sections are measured for events containing a $Z$ boson decaying into electrons or muons and produced in association with at least one or at least two $b$-jets with transverse momentum $p_T> 20$ GeV and rapidity $|y|<2.5$. Predictions from several Monte Carlo generators based on leading-order (LO) or next-to-leading-order (NLO) matrix elements interfaced with a parton-shower simulation and testing different flavour schemes for the choice of initial-state partons are compared with measured cross-sections. The 5-flavour number scheme predictions at NLO accuracy agree better with data than 4-flavour number scheme ones. The 4-flavour number scheme predictions underestimate data in events with at least one $b$-jet. N.B.: Data correspond to average of electron and muon channel. Use LMODE=EL,MU to run on individual channels.