wenmin-wu/ds-skills/2d-to-3d-model-conversion

import torch
import torch.nn as nn
def convert_3d(module):
    module_output = module
    if isinstance(module, nn.BatchNorm2d):
        module_output = nn.BatchNorm3d(module.num_features, module.eps,
                                        module.momentum, module.affine,
                                        module.track_running_stats)
        if module.affine:
            with torch.no_grad():
                module_output.weight = module.weight
                module_output.bias = module.bias
            module_output.running_mean = module.running_mean
            module_output.running_var = module.running_var
    elif isinstance(module, nn.Conv2d):
        module_output = nn.Conv3d(
            module.in_channels, module.out_channels,
            kernel_size=module.kernel_size[0], stride=module.stride[0],
            padding=module.padding[0], dilation=module.dilation[0],
            groups=module.groups, bias=module.bias is not None)
        # Inflate: repeat 2D weights along depth
        module_output.weight = nn.Parameter(
            module.weight.unsqueeze(-1).repeat(1, 1, 1, 1, module.kernel_size[0]))
    elif isinstance(module, nn.MaxPool2d):
        module_output = nn.MaxPool3d(
            kernel_size=module.kernel_size, stride=module.stride,
            padding=module.padding)
    for name, child in module.named_children():
        module_output.add_module(name, convert_3d(child))
    return module_output
# Usage: inflate a pretrained EfficientNet to 3D
model_2d = torchvision.models.efficientnet_v2_s(weights="IMAGENET1K_V1")
model_3d = convert_3d(model_2d)