def __init__(self): super().__init__() factor = 1 base_filters = [64, 128, 256, 512, 512, 2] filters = [f * factor for f in base_filters] self.convs = trw.layers.convs_2d( 6, filters, convolution_kernels=[4, 4, 4, 4, 4, 1], strides=[2, 2, 2, 2, 1, 1], pooling_size=None, dropout_probability=0.2, activation=functools.partial(nn.LeakyReLU, negative_slope=0.1), last_layer_is_output=True, config=default_layer_config(dimensionality=None, norm_type=NormType.InstanceNorm)) self.shapes = [64, 128, 256, 256] self.networks = nn.ModuleList() for s in self.shapes: self.networks.append( trw.layers.convs_2d(6, filters, convolution_kernels=[4, 4, 4, 4, 4, 1], strides=[2, 2, 2, 2, 1, 1], pooling_size=None, dropout_probability=0.2, activation=functools.partial( nn.LeakyReLU, negative_slope=0.1), last_layer_is_output=True, config=default_layer_config( dimensionality=None, norm_type=NormType.InstanceNorm)))
def test_layer_res(self): config = default_layer_config(dimensionality=2) b = BlockRes(config, 8, kernel_size=7, padding='same', padding_mode='reflect') i = torch.zeros([2, 8, 16, 16]) o = b(i) assert o.shape == (2, 8, 16, 16) version = torch.__version__[:3] ops_b = list(b.block_1.ops) assert len(ops_b) == 3 assert isinstance(ops_b[0].ops[0], torch.nn.Conv2d) if version != '1.0': assert ops_b[0].ops[0].padding_mode == 'reflect' assert isinstance(ops_b[1], torch.nn.BatchNorm2d) assert isinstance(ops_b[2], torch.nn.ReLU) ops_b = list(b.block_2.ops) assert len(ops_b) == 2 assert isinstance(ops_b[0].ops[0], torch.nn.Conv2d) if version != '1.0': assert ops_b[0].ops[0].padding_mode == 'reflect' assert isinstance(ops_b[1], torch.nn.BatchNorm2d)
def __init__(self): super().__init__() config = default_layer_config(conv_kwargs={ 'padding': 'same', 'bias': False }, deconv_kwargs={ 'padding': 'same', 'bias': False }) generator = trw.layers.EncoderDecoderResnet( dimensionality=2, input_channels=3, output_channels=3, convolution_kernel=3, encoding_channels=[64, 128, 256], decoding_channels=[256, 128, 64], decoding_block=BlockUpsampleNnConvNormActivation, init_block=functools.partial(BlockConvNormActivation, kernel_size=7), out_block=functools.partial(BlockConvNormActivation, kernel_size=7), config=config) self.generator = generator
def __init__(self): super().__init__() config = default_layer_config(conv_kwargs={ 'padding': 'same', 'bias': True, 'padding_mode': 'reflect' }, deconv_kwargs={ 'padding': 'same', 'bias': True, 'padding_mode': 'reflect' }, norm_type=NormType.InstanceNorm) channels = [32, 64, 128] generator = trw.layers.EncoderDecoderResnet( dimensionality=2, input_channels=3, output_channels=3, encoding_channels=channels, decoding_channels=list(reversed(channels)), decoding_block=BlockUpsampleNnConvNormActivation, init_block=functools.partial(BlockConvNormActivation, kernel_size=7), out_block=functools.partial(BlockConvNormActivation, kernel_size=7), config=config) self.generator = generator
def __init__(self): super().__init__() factor = 1 base_filters = [64, 128, 256, 2] filters = [f * factor for f in base_filters] config = default_layer_config(conv_kwargs={ 'padding': 'same', 'bias': True, 'padding_mode': 'reflect' }, deconv_kwargs={ 'padding': 'same', 'bias': True, 'padding_mode': 'reflect' }, norm_type=NormType.InstanceNorm) self.convs = trw.layers.convs_2d(6, filters, convolution_kernels=[4, 4, 4, 1], strides=[2, 2, 2, 1], pooling_size=None, dropout_probability=0.2, activation=functools.partial( nn.LeakyReLU, negative_slope=0.2), last_layer_is_output=True, config=config)
def test_encoder_decoder_res_k4(self): config = default_layer_config(conv_kwargs={ 'padding': 'same', 'bias': False }, deconv_kwargs={ 'padding': 'same', 'bias': False }) I_O = functools.partial(BlockConvNormActivation, kernel_size=7) model = EncoderDecoderResnet( 2, 3, 2, encoding_channels=[16, 32, 64], decoding_channels=[64, 32, 16], convolution_kernel=4, init_block=I_O, out_block=I_O, config=config, decoding_block=BlockUpsampleNnConvNormActivation, activation=nn.LeakyReLU) t = torch.zeros([5, 3, 32, 64]) o = model(t) assert o.shape == (5, 2, 32, 64)
def test_encoder_decoder_res(self): config = default_layer_config( conv_kwargs={'padding': 'same', 'bias': False}, deconv_kwargs={'padding': 'same', 'bias': False} ) I_O = functools.partial(BlockConvNormActivation, kernel_size=7) model = EncoderDecoderResnet( 2, 3, 2, encoding_channels=[16, 32, 64], decoding_channels=[64, 32, 16], convolution_kernel=5, init_block=I_O, out_block=I_O, config=config, activation=nn.LeakyReLU ) t = torch.zeros([5, 3, 32, 64]) o = model(t) assert o.shape == (5, 2, 32, 64) assert len(model.initial.ops) == 3 assert model.initial.ops[0].kernel_size == (7, 7) assert model.initial.ops[0].bias is None assert isinstance(model.initial.ops[2], nn.LeakyReLU) assert len(model.encoders) == 3 ops = model.encoders[0].ops assert len(ops) == 3 assert ops[0].kernel_size == (5, 5) assert ops[0].bias is None assert isinstance(ops[2], nn.LeakyReLU) ops = model.encoders[1].ops assert len(ops) == 3 assert ops[0].kernel_size == (5, 5) assert ops[0].bias is None assert isinstance(ops[2], nn.LeakyReLU) ops = model.encoders[2].ops assert len(ops) == 3 assert ops[0].kernel_size == (5, 5) assert ops[0].bias is None assert isinstance(ops[2], nn.LeakyReLU) assert len(model.decoders) == 3 ops = model.decoders[2].ops assert len(ops) == 3 assert ops[0].kernel_size == (5, 5) assert ops[0].bias is None assert isinstance(ops[2], nn.LeakyReLU) ops = model.out.ops assert len(ops) == 2 assert ops[0].kernel_size == (7, 7) assert ops[0].bias is None assert isinstance(ops[1], nn.BatchNorm2d)
def __init__(self): super().__init__() config = default_layer_config(dimensionality=2, norm_type=NormType.InstanceNorm) #config = default_layer_config(norm_type=NormType.BatchNorm) """ channels = [32, 64, 128, 256] generator = UpsamplingGenerator( dimensionality=2, input_channels=0, channels=channels + [3], stride=[2] * len(channels) + [1], conditional_channels=3, config=config, last_layer_is_output=True, blocks_with_supervision=[0, 3] ) """ output_channels = 3 channels = [64, 128, 256] intermediates = [128, 64, 32] self.intermediates_outputs = nn.ModuleList() for i in intermediates: config_intermediate = copy.deepcopy(config) config_intermediate.activation = None config_intermediate.norm = None self.intermediates_outputs.append( trw.layers.BlockConvNormActivation( config=config_intermediate, input_channels=i, output_channels=output_channels, kernel_size=1)) """ generator = trw.layers.EncoderDecoderResnet( dimensionality=2, input_channels=3, output_channels=3, convolution_kernel=3, encoding_channels=channels, decoding_channels=list(reversed(channels)), decoding_block=BlockUpsampleNnConvNormActivation, init_block=functools.partial(BlockConvNormActivation, kernel_size=7), out_block=functools.partial(BlockConvNormActivation, kernel_size=7), config=config ) """ generator = trw.layers.UNetBase(dim=2, input_channels=3, channels=channels, output_channels=3, config=config) self.generator = generator self.scale = 2**len(channels)
def __init__( self, dimensionality: int, input_channels: int, encoder_channels: Sequence[int], decoder_channels: Sequence[int], convolution_kernels: Optional[Union[int, Sequence[int]]] = 5, encoder_strides: Union[int, List[int], NestedIntSequence] = 1, decoder_strides: Union[int, List[int], NestedIntSequence] = 2, pooling_size: Optional[Union[int, List[int], NestedIntSequence]] = 2, convolution_repeats: Union[int, Sequence[int]] = 1, activation: Optional[Activation] = nn.ReLU, dropout_probability: Optional[float] = None, norm_type: NormType = NormType.BatchNorm, norm_kwargs: Dict = {}, activation_kwargs: Dict = {}, last_layer_is_output: bool = False, force_decoded_size_same_as_input: bool = True, squash_function: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, config: LayerConfig = default_layer_config(dimensionality=None)): super().__init__() self.force_decoded_size_same_as_input = force_decoded_size_same_as_input self.encoder = ConvsBase(dimensionality=dimensionality, input_channels=input_channels, channels=encoder_channels, convolution_kernels=convolution_kernels, strides=encoder_strides, pooling_size=pooling_size, convolution_repeats=convolution_repeats, activation=activation, dropout_probability=dropout_probability, norm_type=norm_type, activation_kwargs=activation_kwargs, norm_kwargs=norm_kwargs, last_layer_is_output=False, config=config) self.decoder = ConvsTransposeBase( dimensionality=dimensionality, input_channels=encoder_channels[-1], channels=decoder_channels, strides=decoder_strides, activation=activation, dropout_probability=dropout_probability, norm_type=norm_type, norm_kwargs=norm_kwargs, activation_kwargs=activation_kwargs, last_layer_is_output=last_layer_is_output, squash_function=squash_function, config=config)
def test_efficient_net_construction(self): config = default_layer_config(dimensionality=2, activation=Swish) torch.random.manual_seed(0) i = torch.randn([4, 3, 224, 224]) net = EfficientNet(dimensionality=2, input_channels=3, output_channels=10, config=config) o = net.feature_extractor(i) assert o.shape == (4, 1280, 7, 7) o = net(i) assert o.shape == (4, 10) intermediates = net.forward_with_intermediate(i) assert len(intermediates) == 9
def test_upsample_nn(self): config = default_layer_config(dimensionality=3, norm_type=NormType.InstanceNorm) block = BlockUpsampleNnConvNormActivation(config, 2, 4, kernel_size=3, stride=(1, 2, 3)) i = torch.zeros([10, 2, 4, 5, 6], dtype=torch.float32) o = block(i) assert o.shape == (10, 4, 4 * 1, 5 * 2, 6 * 3) ops = list(block.ops) assert len(ops) == 4 assert isinstance(ops[0], nn.Upsample) assert isinstance(ops[1], nn.Conv3d) assert isinstance(ops[2], nn.InstanceNorm3d) assert isinstance(ops[3], nn.ReLU)
def __init__(self, options): super().__init__() config = default_layer_config() I_O = functools.partial(BlockConvNormActivation, kernel_size=7) self.model = trw.layers.EncoderDecoderResnet( 3, input_channels=1, output_channels=3, encoding_channels=[64, 128], decoding_channels=[128, 64], init_block=I_O, out_block=I_O, config=config, nb_residual_blocks=18, ) self.norm_input = nn.InstanceNorm3d(1) self.options = options
def __init__(self): super().__init__() self.init_block = BlockConvNormActivation(config=default_layer_config( norm_type=None, dimensionality=2, activation=functools.partial(nn.LeakyReLU, negative_slope=0.2), ), input_channels=6, output_channels=64, kernel_size=4, stride=2) self.convs = trw.layers.convs_2d(input_channels=64, channels=[128, 256, 512, 2], convolution_kernels=4, strides=[2, 2, 1, 1], pooling_size=None, last_layer_is_output=True, activation=functools.partial( nn.LeakyReLU, negative_slope=0.2))
def __init__( self, unet: UNetBase, input_target_shape: ShapeCX, output_creator: Callable[[TorchTensorNCX, TensorNCX], Output], output_block: ConvBlockType = BlockConvNormActivation, discard_top_k_outputs: int = 1, resize_mode: Literal['nearest', 'linear'] = 'nearest', config: LayerConfig = default_layer_config(dimensionality=None)): super().__init__() self.unet = unet assert len( input_target_shape ) == unet.dim + 1, 'must be a shape with `N` component removed!' device = get_device(unet) self.discard_top_k_outputs = discard_top_k_outputs # dummy test to get the intermediate layer shapes dummy_input = torch.zeros([1] + list(input_target_shape), device=device) outputs = unet.forward_with_intermediate(dummy_input) # no activation, these are all output nodes! config = copy.copy(config) config.set_dim(unet.dim) config.activation = None config.norm = None self.outputs = nn.ModuleList() for o in outputs[discard_top_k_outputs:]: output = output_block(config, o.shape[1], self.unet.output_channels) self.outputs.append(output) self.output_creator = output_creator self.resize_mode = resize_mode
def test_conv2(self): net = trw.layers.ConvsBase( 3, 2, channels=[4, 8, 10], convolution_kernels=3, strides=2, pooling_size=2, norm_type=trw.layers.NormType.InstanceNorm, convolution_repeats=2, activation=nn.LeakyReLU, with_flatten=True, config=default_layer_config(dimensionality=None), last_layer_is_output=True) # # First Block # assert len(net.layers) == 3 children_0 = list(net.layers[0].children()) assert len(children_0) == 3 children_00 = list(net.layers[0][0].ops) assert len(children_00) == 3 assert isinstance(children_00[0], nn.Conv3d) assert children_00[0].weight.shape[1] == 2 assert children_00[0].stride == (1, 1, 1) assert isinstance(children_00[1], nn.InstanceNorm3d) assert isinstance(children_00[2], nn.LeakyReLU) children_01 = list(net.layers[0][1].ops) assert len(children_01) == 3 assert isinstance(children_01[0], nn.Conv3d) assert children_01[0].weight.shape[1] == 4 assert children_01[0].stride == (2, 2, 2) assert isinstance(children_01[1], nn.InstanceNorm3d) assert isinstance(children_01[2], nn.LeakyReLU) children_02 = net.layers[0][2].op assert isinstance(children_02, nn.MaxPool3d) # # Second block # children_1 = list(net.layers[1].children()) assert len(children_1) == 3 children_10 = list(net.layers[1][0].ops) assert len(children_10) == 3 assert isinstance(children_10[0], nn.Conv3d) assert children_10[0].weight.shape[1] == 4 assert children_10[0].stride == (1, 1, 1) assert isinstance(children_10[1], nn.InstanceNorm3d) assert isinstance(children_10[2], nn.LeakyReLU) children_11 = list(net.layers[1][1].ops) assert len(children_11) == 3 assert isinstance(children_11[0], nn.Conv3d) assert children_11[0].weight.shape[1] == 8 assert children_11[0].stride == (2, 2, 2) assert isinstance(children_11[1], nn.InstanceNorm3d) assert isinstance(children_11[2], nn.LeakyReLU) children_12 = net.layers[1][2].op assert isinstance(children_12, nn.MaxPool3d) # # Third block # children_2 = list(net.layers[2].children()) assert len(children_2) == 3 children_20 = list(net.layers[2][0].ops) assert len(children_20) == 3 assert isinstance(children_20[0], nn.Conv3d) assert children_20[0].stride == (1, 1, 1) assert isinstance(children_20[1], nn.InstanceNorm3d) assert isinstance(children_20[2], nn.LeakyReLU) children_21 = list(net.layers[2][1].ops) assert len(children_21) == 1 assert isinstance(children_21[0], nn.Conv3d) assert children_11[0].stride == (2, 2, 2) children_22 = net.layers[2][2].op assert isinstance(children_22, nn.MaxPool3d)
def __init__( self, dimensionality: int, input_channels: int, channels: Sequence[int], *, convolution_kernels: Union[int, List[int], NestedIntSequence] = 5, strides: Union[int, List[int], NestedIntSequence] = 2, paddings: Optional[Union[str, int, List[int], NestedIntSequence]] = None, activation: Any = nn.ReLU, activation_kwargs: Dict = {}, dropout_probability: Optional[float] = None, norm_type: Optional[NormType] = None, norm_kwargs: Dict = {}, last_layer_is_output: bool = False, squash_function: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, deconv_block_fn: ConvTransposeBlockType = BlockDeconvNormActivation, config: LayerConfig = default_layer_config(dimensionality=None), target_shape: Optional[Sequence[int]] = None): """ Args: dimensionality: the dimension of the CNN (2 for 2D or 3 for 3D) input_channels: the number of input channels channels: the number of channels for each convolutional layer convolution_kernels: for each convolution group, the kernel of the convolution strides: for each convolution group, the stride of the convolution dropout_probability: if None, not dropout. Else the probability of dropout after each convolution norm_kwargs: the normalization additional arguments. See the original torch functions for description. last_layer_is_output: if True, the last convolution will NOT have activation, dropout, batch norm, LRN squash_function: a function to be applied on the reconstuction. It is common to apply for example ``torch.sigmoid``. If ``None``, no function applied paddings: the paddings added. If ``None``, half the convolution kernel will be used. target_shape: if not ``None``, the output layer will be cropped or padded to mach the target (N, C components excluded) """ super().__init__() # update the configuration locally config = copy.copy(config) if norm_type is not None: config.norm_type = norm_type if activation is not None: config.activation = activation config.set_dim(dimensionality) config.norm_kwargs = {**norm_kwargs, **config.activation_kwargs} config.activation_kwargs = { **activation_kwargs, **config.activation_kwargs } # normalize the arguments nb_convs = len(channels) if not isinstance(convolution_kernels, list): convolution_kernels = [convolution_kernels] * nb_convs if not isinstance(strides, list): strides = [strides] * nb_convs if paddings is None: paddings = [div_shape(kernel, 2) for kernel in convolution_kernels] elif isinstance(paddings, numbers.Integral): paddings = [paddings] * nb_convs else: assert isinstance( paddings, collections.Sequence) and len(paddings) == nb_convs assert nb_convs == len( convolution_kernels ), 'must be specified for each convolutional layer' assert nb_convs == len( strides), 'must be specified for each convolutional layer' layers = nn.ModuleList() prev = input_channels for n in range(len(channels)): current = channels[n] currently_last_layer = n + 1 == len(channels) p = paddings[n] if last_layer_is_output and currently_last_layer: # Last layer layer should not have dropout/normalization/activation config.norm = None config.activation = None config.dropout = None ops = [] ops.append( deconv_block_fn( config, prev, current, kernel_size=convolution_kernels[n], stride=strides[n], padding=p, output_padding=strides[n] - 1, )) if config.dropout is not None and dropout_probability is not None: ops.append( config.dropout(p=dropout_probability, **config.dropout_kwargs)) layers.append(nn.Sequential(*ops)) prev = current self.layers = layers self.squash_function = squash_function self.target_shape = target_shape
def __init__(self, dimensionality: int, input_channels: int, channels: Sequence[int], conditional_channels: int, blocks_with_supervision: Sequence[int], *, kernel_size: ConvKernels = 3, stride: Optional[Stride] = 2, resize_mode: Literal['nearest', 'linear'] = 'nearest', upsampling_block: ConvTransposeBlockType = BlockUpsampleNnConvNormActivation, last_layer_is_output: bool = True, config: LayerConfig = default_layer_config()): super().__init__() self.dimensionality = dimensionality self.channels = channels self.input_channels = input_channels self.resize_mode = resize_mode self.conditional_channels = conditional_channels assert len(blocks_with_supervision) >= 1 config = copy.deepcopy(config) config.set_dim(dimensionality) if stride is not None: if isinstance(stride, list): assert len(stride) == len(channels) else: stride = [stride] * len(channels) if kernel_size is not None: config.conv_kwargs['kernel_size'] = kernel_size config.deconv_kwargs['kernel_size'] = kernel_size self.blocks = nn.ModuleList() self.blocks_with_supervision = sorted(blocks_with_supervision) supervision_index = 0 i = input_channels for block_index, o in enumerate(channels): if supervision_index < len(self.blocks_with_supervision) and \ block_index == self.blocks_with_supervision[supervision_index]: # we will receive additional supervision here. Increase the # channels by the number of input channels. Supervision MUST be a # down-scaled version of the full scale conditional image i = i + self.conditional_channels supervision_index += 1 if stride is not None: config.deconv_kwargs['stride'] = stride[block_index] if last_layer_is_output and block_index + 1 == len(channels): # last block: remove activation and normalization layers config.activation = None config.norm = None block = upsampling_block(config=config, input_channels=i, output_channels=o) i = o self.blocks.append(block)