model.py 文件源码

def __init__(self,
                 output_ch, # outputs channel, size is same as inputs
                 block_fn = 'origin',
                 name = 'unet'):
        self.output_ch = output_ch
        self.name = name

        assert block_fn in ['batch_norm', 'origin'], 'choose \'batch_norm\' or \'origin\''
        if block_fn == 'batch_norm':
            self.block_fn = _conv_bn_relu
        elif block_fn == 'origin':
            self.block_fn = _conv_relu