Function save_unet_model

python/unet.py:144–152

(model, path)

    save_res_transformer(output_blocks[11], pathlib.Path(path, 'rt7'))

def save_unet_model(model, path):
    pathlib.Path(path).mkdir(parents=True, exist_ok=True)
    save_linear(model.time_embed[0], pathlib.Path(path, 'lin1_time_embed'))
    save_linear(model.time_embed[2], pathlib.Path(path, 'lin2_time_embed'))
    save_unet_input_blocks(model.input_blocks, pathlib.Path(path, 'input_blocks'))
    save_res_transformer_res(model.middle_block, pathlib.Path(path, 'middle_block'))
    save_unet_output_blocks(model.output_blocks, pathlib.Path(path, 'output_blocks'))
    save_group_norm(model.out[0], pathlib.Path(path, 'norm_out'))
    save_conv2d(model.out[2], pathlib.Path(path, 'conv_out'))