Function save_res_block

python/unet.py:8–30

(res_block, path)

from tinygrad.nn import Conv2d

def save_res_block(res_block, path):
    pathlib.Path(path).mkdir(parents=True, exist_ok=True)
    # We can't directly save activation functions, but as they are just attribute of the block,
    # we don't need to save them separately, they will be recreated along with the block.
    
    # saving group normalization layer
    save_group_norm(res_block.in_layers[0], os.path.join(path, 'norm_in'))
    
    # saving the convolutional layer
    save_conv2d(res_block.in_layers[2], os.path.join(path, 'conv_in'))

    # saving the linear layer
    save_linear(res_block.emb_layers[1], os.path.join(path, 'lin_embed'))
    
    # saving group normalization in out_layers
    save_group_norm(res_block.out_layers[0], os.path.join(path, 'norm_out'))

    # saving the convolutional layer in out_layers
    save_conv2d(res_block.out_layers[3], os.path.join(path, 'conv_out'))

    # save skip_connection based on the object type
    if isinstance(res_block.skip_connection, Conv2d):
        save_conv2d(res_block.skip_connection, os.path.join(path, 'skip_connection'))

def save_cross_attention(cross_attention, path):
    pathlib.Path(path).mkdir(parents=True, exist_ok=True)