def transpose_checkpoint_model_state_nhwc_to_nchw(model_dict):
for k in model_dict:
param_tensor = model_dict[k]
needs_transpose = is_layer_nhwc_eligible(k) and len(
param_tensor.shape) == 4
if needs_transpose:
model_dict[k] = model_dict[k].permute(0, 3, 1, 2).contiguous()
def transpose_optimizer_state_nchw_to_nhwc(model, optimizer_dict):
layer_id_to_name_map = {}
for name, param in model.named_parameters():
layer_id_to_name_map[id(param)] = name
for k in optimizer_dict['state']:
needs_transpose = is_layer_nhwc_eligible(layer_id_to_name_map[k])
needs_transpose = needs_transpose and \
len(optimizer_dict['state'][k]['momentum_buffer'].shape) == 4
if needs_transpose:
optimizer_dict['state'][k]['momentum_buffer'] = \
optimizer_dict['state'][k]['momentum_buffer'].permute(0,2,3,1).contiguous()