fepegar / unet

"pip install unet": PyTorch Implementation of 1D, 2D and 3D U-Net architecture.
MIT License
148 stars 22 forks source link

Different kernel sizes for different dimensions #40

Open kuangdai opened 6 months ago

kuangdai commented 6 months ago

Can you maybe add this feature, please?

I think the main change would be

class ConvolutionalBlock(nn.Module):
    def __init__(
            self,
            dimensions: int,
            in_channels: int,
            out_channels: int,
            normalization: Optional[str] = None,
            kernel_size: int = 3,
            activation: Optional[str] = 'ReLU',
            preactivation: bool = False,
            padding: int = 0,
            padding_mode: str = 'zeros',
            dilation: Optional[int] = None,
            dropout: float = 0,
            ):
class ConvolutionalBlock(nn.Module):
    def __init__(
            self,
            dimensions: int,
            in_channels: int,
            out_channels: int,
            normalization: Optional[str] = None,
            kernel_size: int = Union[int, Sequence[int]],
            activation: Optional[str] = 'ReLU',
            preactivation: bool = False,
            padding: int = 0,
            padding_mode: str = 'zeros',
            dilation: Optional[int] = None,
            dropout: float = 0,
            ):

I can create a PR if you prefer this way.

fepegar commented 6 months ago

Hi, @kuangdai. Please feel free to open the PR!

lizhuoq commented 4 months ago

If kernel_size can accept a tuple, doesn't it mean that padding should also accept a tuple?