tensorflow2.7使用functional api实现unet

from tensorflow.keras.layers import (Input, Conv2D, MaxPooling2D,
                                     UpSampling2D,Concatenate)
from tensorflow.keras.models import Model


img_rows = 512
img_cols = 512

def UNET():
    inputs = Input(shape=(img_rows,img_cols,1),name='initial_input')
    conv1 = Conv2D(32, 3, activation='relu', padding='same',name='NO1_conv_1')(inputs)
    conv1 = Conv2D(32, 3, activation='relu', padding='same',name='NO1_conv_2')(conv1)
    pool1 = MaxPooling2D(name='NO1_maxpooling')(conv1)

    conv2 = Conv2D(64, 3, activation='relu', padding='same',name='NO2_conv_1')(pool1)
    conv2 = Conv2D(64, 3, activation='relu', padding='same',name='NO2_conv_2')(conv2)
    pool2 = MaxPooling2D(name='NO2_maxpooling')(conv2)

    conv3 = Conv2D(128, 3, activation='relu', padding='same',name='NO3_conv_1')(pool2)
    conv3 = Conv2D(128, 3, activation='relu', padding='same',name='NO3_conv_2')(conv3)
    pool3 = MaxPooling2D(name='NO3_maxpooling')(conv3)

    conv4 = Conv2D(256, 3, activation='relu', padding='same',name='NO4_conv_1')(pool3)
    conv4 = Conv2D(256, 3, activation='relu', padding='same',name='NO4_conv_2')(conv4)
    pool4 = MaxPooling2D(name='NO4_maxpooling')(conv4)

    conv5 = Conv2D(512, 3, activation='relu', padding='same',name='NO5_conv_1')(pool4)
    conv5 = Conv2D(512, 3, activation='relu', padding='same',name='NO5_conv_2')(conv5)
    conv5_up = UpSampling2D(name='conv5_up')(conv5)

    up6 = Concatenate(axis=3,name='up6')([conv5_up,conv4])
    conv6 = Conv2D(256, 3, activation='relu', padding='same',name='NO6_conv_1')(up6)
    conv6 = Conv2D(256, 3, activation='relu', padding='same',name='NO6_conv_2')(conv6)
    conv6_up = UpSampling2D(name='conv6_up')(conv6)

    up7 = Concatenate(axis=3,name='up7')([conv6_up,conv3])
    conv7 = Conv2D(128,3, activation='relu', padding='same',name='NO7_conv_1')(up7)
    conv7 = Conv2D(128,3, activation='relu', padding='same',name='NO7_conv_2')(conv7)
    conv7_up = UpSampling2D(name='conv7_up')(conv7)

    up8 = Concatenate(axis=3,name='up8')([conv7_up,conv2])
    conv8 = Conv2D(64,3, activation='relu', padding='same',name='NO8_conv_1')(up8)
    conv8 = Conv2D(64,3, activation='relu', padding='same',name='NO8_conv_2')(conv8)
    conv8_up = UpSampling2D(name='conv8_up')(conv8)

    up9 = Concatenate(axis=3,name='up9')([conv8_up,conv1])
    conv9 = Conv2D(32, 3, activation='relu', padding='same',name='NO9_conv_1')(up9)
    conv9 = Conv2D(32, 3, activation='relu', padding='same',name='NO9_conv_2')(conv9)

    conv10 = Conv2D(1, 1, activation='sigmoid',name='cls_layer')(conv9) # num_classes=1, 1*1卷积,相当于dense层

    model = Model(inputs=inputs, outputs=conv10,name='my_unet_tf2')
    model.summary()


if __name__ == '__main__':
    UNET()

 

Model: "my_unet_tf2"
__________________________________________________________________________________________________
 Layer (type)                   Output Shape         Param #     Connected to                     
==================================================================================================
 initial_input (InputLayer)     [(None, 512, 512, 1  0           []                               
                                )]                                                                
                                                                                                  
 NO1_conv_1 (Conv2D)            (None, 512, 512, 32  320         ['initial_input[0][0]']          
                                )                                                                 
                                                                                                  
 NO1_conv_2 (Conv2D)            (None, 512, 512, 32  9248        ['NO1_conv_1[0][0]']             
                                )                                                                 
                                                                                                  
 NO1_maxpooling (MaxPooling2D)  (None, 256, 256, 32  0           ['NO1_conv_2[0][0]']             
                                )                                                                 
                                                                                                  
 NO2_conv_1 (Conv2D)            (None, 256, 256, 64  18496       ['NO1_maxpooling[0][0]']         
                                )                                                                 
                                                                                                  
 NO2_conv_2 (Conv2D)            (None, 256, 256, 64  36928       ['NO2_conv_1[0][0]']             
                                )                                                                 
                                                                                                  
 NO2_maxpooling (MaxPooling2D)  (None, 128, 128, 64  0           ['NO2_conv_2[0][0]']             
                                )                                                                 
                                                                                                  
 NO3_conv_1 (Conv2D)            (None, 128, 128, 12  73856       ['NO2_maxpooling[0][0]']         
                                8)                                                                
                                                                                                  
 NO3_conv_2 (Conv2D)            (None, 128, 128, 12  147584      ['NO3_conv_1[0][0]']             
                                8)                                                                
                                                                                                  
 NO3_maxpooling (MaxPooling2D)  (None, 64, 64, 128)  0           ['NO3_conv_2[0][0]']             
                                                                                                  
 NO4_conv_1 (Conv2D)            (None, 64, 64, 256)  295168      ['NO3_maxpooling[0][0]']         
                                                                                                  
 NO4_conv_2 (Conv2D)            (None, 64, 64, 256)  590080      ['NO4_conv_1[0][0]']             
                                                                                                  
 NO4_maxpooling (MaxPooling2D)  (None, 32, 32, 256)  0           ['NO4_conv_2[0][0]']             
                                                                                                  
 NO5_conv_1 (Conv2D)            (None, 32, 32, 512)  1180160     ['NO4_maxpooling[0][0]']         
                                                                                                  
 NO5_conv_2 (Conv2D)            (None, 32, 32, 512)  2359808     ['NO5_conv_1[0][0]']             
                                                                                                  
 conv5_up (UpSampling2D)        (None, 64, 64, 512)  0           ['NO5_conv_2[0][0]']             
                                                                                                  
 up6 (Concatenate)              (None, 64, 64, 768)  0           ['conv5_up[0][0]',               
                                                                  'NO4_conv_2[0][0]']             
                                                                                                  
 NO6_conv_1 (Conv2D)            (None, 64, 64, 256)  1769728     ['up6[0][0]']                    
                                                                                                  
 NO6_conv_2 (Conv2D)            (None, 64, 64, 256)  590080      ['NO6_conv_1[0][0]']             
                                                                                                  
 conv6_up (UpSampling2D)        (None, 128, 128, 25  0           ['NO6_conv_2[0][0]']             
                                6)                                                                
                                                                                                  
 up7 (Concatenate)              (None, 128, 128, 38  0           ['conv6_up[0][0]',               
                                4)                                'NO3_conv_2[0][0]']             
                                                                                                  
 NO7_conv_1 (Conv2D)            (None, 128, 128, 12  442496      ['up7[0][0]']                    
                                8)                                                                
                                                                                                  
 NO7_conv_2 (Conv2D)            (None, 128, 128, 12  147584      ['NO7_conv_1[0][0]']             
                                8)                                                                
                                                                                                  
 conv7_up (UpSampling2D)        (None, 256, 256, 12  0           ['NO7_conv_2[0][0]']             
                                8)                                                                
                                                                                                  
 up8 (Concatenate)              (None, 256, 256, 19  0           ['conv7_up[0][0]',               
                                2)                                'NO2_conv_2[0][0]']             
                                                                                                  
 NO8_conv_1 (Conv2D)            (None, 256, 256, 64  110656      ['up8[0][0]']                    
                                )                                                                 
                                                                                                  
 NO8_conv_2 (Conv2D)            (None, 256, 256, 64  36928       ['NO8_conv_1[0][0]']             
                                )                                                                 
                                                                                                  
 conv8_up (UpSampling2D)        (None, 512, 512, 64  0           ['NO8_conv_2[0][0]']             
                                )                                                                 
                                                                                                  
 up9 (Concatenate)              (None, 512, 512, 96  0           ['conv8_up[0][0]',               
                                )                                 'NO1_conv_2[0][0]']             
                                                                                                  
 NO9_conv_1 (Conv2D)            (None, 512, 512, 32  27680       ['up9[0][0]']                    
                                )                                                                 
                                                                                                  
 NO9_conv_2 (Conv2D)            (None, 512, 512, 32  9248        ['NO9_conv_1[0][0]']             
                                )                                                                 
                                                                                                  
 cls_layer (Conv2D)             (None, 512, 512, 1)  33          ['NO9_conv_2[0][0]']             
                                                                                                  
==================================================================================================
Total params: 7,846,081
Trainable params: 7,846,081
Non-trainable params: 0
__________________________________________________________________________________________________

Process finished with exit code 0