Python: Memory consumption exceeds what's expected

In my code below, I have tried profiling the memory usage of one function.

Line #    Mem usage    Increment   Line Contents
================================================
149    100.6 MiB      0.0 MiB   @profile
150                             def main():
151    100.6 MiB      0.0 MiB       print('Opening CSV file...')
152    107.1 MiB      6.5 MiB       df = open_csv_file()
153    107.1 MiB      0.0 MiB       print('Loading feature array...')
154    288.4 MiB    181.3 MiB       graph_array = load_feat_array()
155    288.4 MiB      0.0 MiB       print('Opening graph_idxs...')
156    314.4 MiB     26.0 MiB       graph_idxs = unpickle_data('../data/graph_idxs.pkl')
157    314.4 MiB      0.0 MiB       print('Opening graph_nodes...')
158    409.7 MiB     95.3 MiB       graph_nodes = unpickle_data('../data/graph_nodes.pkl')
159    409.7 MiB      0.0 MiB       print('Opening nodes_nbrs...')
160    638.6 MiB    228.9 MiB       nodes_nbrs = unpickle_data('../data/nodes_nbrs.pkl')
161
162                                 # Check data
163    638.6 MiB      0.0 MiB       print('Doing data checks...')
164    638.6 MiB      0.0 MiB       assert df.shape == (6660, 13)
165    638.6 MiB      0.0 MiB       assert len(graph_array) == 659895
166    638.6 MiB      0.0 MiB       assert len(graph_idxs) == len(graph_nodes)
167    638.6 MiB      0.0 MiB       assert len(nodes_nbrs) == len(graph_array)
168
169    638.6 MiB      0.0 MiB       print('Preprocessing data...')
170    638.6 MiB      0.0 MiB       pp_data = preprocess_data(df, nodes_nbrs, graph_idxs, graph_nodes,
171   1022.8 MiB    384.2 MiB                                 graph_array)
172    808.2 MiB   -214.7 MiB       graph_arr, nodes_nbrs, graph_idxs, graph_nodes = pp_data
173
174    808.2 MiB      0.0 MiB       assert graph_arr.shape[0] == len(nodes_nbrs)
175    808.2 MiB      0.0 MiB       assert len(graph_idxs) == len(graph_nodes)
176
177    808.2 MiB      0.0 MiB       print('Setting up neural net.')
178    808.2 MiB      0.0 MiB       layers = [GraphConvLayer(weights_shape=(36, 36),
179    808.2 MiB      0.0 MiB                                biases_shape=(1, 36)),
180    808.2 MiB      0.0 MiB                 FingerprintLayer(weights_shape=(36, 36),
181    808.2 MiB      0.0 MiB                                  biases_shape=(1, 36)),
182    808.2 MiB      0.0 MiB                 LinearRegressionLayer(weights_shape=(36, 1),
183    808.2 MiB      0.0 MiB                                       biases_shape=(1, 1)),
184                                           ]
185    808.2 MiB      0.0 MiB       print(layers)
186
187    808.2 MiB      0.0 MiB       print('Initializing network...')
188    808.2 MiB      0.0 MiB       wb = initialize_network(layers_spec=layers)
189    808.2 MiB      0.0 MiB       wb_vect, unflattener = flatten(wb)
190    808.2 MiB      0.0 MiB       print('Network initialized. Weights & biases:')
191    808.2 MiB      0.0 MiB       print(wb)
192
193    896.3 MiB     88.1 MiB       node_rows, node_cols, ones = to_sparse_format(nodes_nbrs)
194
195    896.3 MiB      0.0 MiB       nodes_nbrs_compressed = csr_matrix((ones, (node_rows, node_cols)),
196    896.3 MiB      0.0 MiB                                          shape=(len(nodes_nbrs),
197    907.8 MiB     11.5 MiB                                                 len(nodes_nbrs)))
198
199    907.8 MiB      0.0 MiB       train_losses = []
200    907.8 MiB      0.0 MiB       preds_iter = []
201    907.8 MiB      0.0 MiB       actual_iter = []
202
203    907.8 MiB      0.0 MiB       print('Defining train loss function.')
204
205   4153.6 MiB   3245.8 MiB       def train_loss(wb_vect, unflattener):
206   4153.6 MiB      0.0 MiB           wb_struct = unflattener(wb_vect)
207   4153.6 MiB      0.0 MiB           preds = predict(wb_struct, graph_arr, nodes_nbrs_compressed,
208   4251.2 MiB     97.6 MiB                           graph_idxs, layers)
209   4252.1 MiB      0.9 MiB           graph_scores = get_actual(graph_idxs, df, preds)
210   4252.1 MiB      0.0 MiB           mse = np.mean(np.power(preds - graph_scores, 2))
211
212   4252.1 MiB      0.0 MiB           train_losses.append(mse)
213   4252.1 MiB      0.0 MiB           preds_iter.append(preds)
214   4252.1 MiB      0.0 MiB           actual_iter.append(graph_scores)
215   4252.1 MiB      0.0 MiB           return mse
216
217    907.8 MiB  -3344.3 MiB       traingrad = grad(train_loss)
218
219    907.8 MiB      0.0 MiB       training_losses = []
220
221    907.8 MiB      0.0 MiB       print('Defining callback function...')
222
223   4153.6 MiB   3245.8 MiB       def callback(wb, i):
224   4153.6 MiB      0.0 MiB           start = time()
225   3385.4 MiB   -768.2 MiB           tl = train_loss(*flatten(wb))
226   3385.4 MiB      0.0 MiB           if i % 1 == 0:
227   3385.4 MiB      0.0 MiB               print(tl, time() - start)
228   3385.4 MiB      0.0 MiB           training_losses.append(tl)
229   3385.4 MiB      0.0 MiB           gc.collect()
230
231    907.8 MiB  -2477.6 MiB       print('Training neural network.')
232   3191.1 MiB   2283.3 MiB       wb_vect, unflattener = adam(traingrad, wb, callback=callback, num_iters=10)
233
234   3191.1 MiB      0.0 MiB       print(wb_vect)

(For background, this is my own implementation of a deep learning thing.)

If you take a look at the memory profiling report, there's a spike in memory usage at line 205, and at line 232. Line 232 is where I call my optimizer.

What baffles me right now is that even with this profile on hand my maximum memory usage (as indicated by the macOS System Profiler) will climb by gigabytes at a time, depending on the number of iterations I take; more iterations = more RAM consumption, with this particular run showing up to 14GB (10GB compressed) usage.

Is there a memory leak here? I can't seem to figure out how to begin debugging this. Would there be any experts who could give some guidance on this?

In case any other details help:

1. I am using autograd to automatically differentiate my code.
2. I have an alternate implementation that I use as a reference, and it does not consume huge amounts of RAM; stays within ~1 GB.
3. I am happy to append the two implementations of underlying code that make the difference between these two, if it helps.