python - 可逆字典到数组映射

Question

我有一个变量名字典，我需要将其传递给优化器，获取输出，然后再设置回字典。字典中的变量有各种形状，但我需要将它们作为一维数组传递。

向前一步很容易:

np.concatenate([elem.ravel() for elem in param_dict.values()])

但是，我不确定将修改后的参数返回到我的字典中的最佳方法是什么。

我知道我可以在巨大的参数数组中为不同的 dict 元素手动找到偏移量并手动将它们 reshape 为正确的形状，但我想知道是否有更简单的方法来做到这一点(例如字典中的元素/形状可以是动态的)。

Answer 1

您可以像这样(滥用)使用结构化数组(或者更确切地说是底层复合数据类型):

# create example 
param_dict = {f'a{i}':np.random.randint(0,10,np.random.randint(1,10,np.random.randint(1,3))) for i in range(1000)}

# make structured array containing a single element
x = np.array((*param_dict.values(),), [(k,(v.dtype,v.shape)) for k,v in param_dict.items()])

# can access elements by key:
x['a7']
# array([[3, 0, 9, 2, 5],
#        [4, 7, 2, 7, 6]])
param_dict['a7']
# array([[3, 0, 9, 2, 5],
#        [4, 7, 2, 7, 6]])

# data are stored flat
import numpy.lib.recfunctions as nlr

flat_view = nlr.structured_to_unstructured(x)
# or simpler:
# flat_view = x[None].view(int)

np.shares_memory(flat_view,x)
# True
flat_view.shape
# (15204,)

# all meta data are preserved in the dtype
x.dtype.fields['a7']
#         type   shape     offset in bytes
# (dtype(('<i8', (2, 5))), 840)

flat_view[840//8:840//8+2*5]
# array([3, 0, 9, 2, 5, 4, 7, 2, 7, 6])

现在让我们将兼容的平面阵列恢复为原始格式

# some token processing
processed = flat_view**2

new_x = nlr.unstructured_to_structured(processed,x.dtype)
# or simpler:
# new_x = processed.view(x.dtype)[0,...]

new_param_dict = {k:new_x[k] for k in new_x.dtype.fields}

new_param_dict['a7']
# array([[ 9,  0, 81,  4, 25],
#        [16, 49,  4, 49, 36]])