How to save dictionary with state_dict to json

1

I ran training on with all permutations of hyper parameters at multiple levels (batch_size, hidden_size, learn_rate). I saved the result of these training’s into a dictionary as follows:

modelList, accList, targList, predList = kFoldCross(model = model[str(params)], 
                                                    data = data, 
                                                    batch_size=batch_size, 
                                                    num_folds=5, 
                                                    device=device,
                                                    train_func = train_modelEdit,
                                                    learn_rate= learn_rate,
                                                    num_epoch= 200,
                                                    patience = 10)

 #make dictionary to store outputs
    bigDict[str(params)] = {'state_dict': modelList, 
                            'accuracy': accList, 
                            'target': targList, 
                            'prediction': predList}

where params is tuple of the hyper parameters i.e. (64, 32, .01) (batch size 64, hidden_size 32, learn_rate .01). modelList is a list where each entry contains the state_dict of the model at a single run of the kfold validitation(5 state_dicts in total). The other lists (accList, targList, and predList) contain ouputs from each of the 5 runs of the kfold.

I tried to save bigDict to json as follows:

with open('bigDict_ConvLSTM_incomplete.json', 'w') as fp:
    json.dump(bigDict, fp)

but I get the following error:

TypeError: Object of type Tensor is not JSON serializable

which presumably refers to the torch.tensors in the state_dict.

Is there a way I can save the entire dictionary to json or do I have to save the model state_dict separately?

In the event that bigDict cannot be saved:
I know I could save the state_dicts individually using torch.save(), but I do not want to have a bunch of different files. Is there at least a way to save them all together?

note: I am open to using any other format not necessarily json to save my files. Any advice is greatly appreciated.

Thank you

2

I would try to store everything in a binary format to avoid the serialization into the (human-readable) JSON format, e.g. by using pickle or torch.save on the dict directly.