tensorflow2.7实现unet网络结构

之前的文章我觉得并没有如实按着论文中的结构来写模型，所以重新写了一下，并在合适的位置标记了代码中的对应位置。

 

unet.py

import warnings
warnings.filterwarnings("ignore")
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

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


def unet():
    """ set up unet model """
    inputs = Input(shape=(572,572,1),name='input')
    conv1_1 = Conv2D(64,3,activation='relu',name='conv1_1')(inputs)
    conv1_2 = Conv2D(64,3,activation='relu',name='conv1_2')(conv1_1)

    max_pool1 = MaxPooling2D(name='maxpool1')(conv1_2) # maxpooling2D默认stride=filer size
    conv2_1 = Conv2D(128,3,activation='relu',name='conv2_1')(max_pool1)
    conv2_2 = Conv2D(128,3,activation='relu',name='conv2_2')(conv2_1)

    max_pool2 = MaxPooling2D(name='maxpool2')(conv2_2)
    conv3_1 = Conv2D(256,3,activation='relu',name='conv3_1')(max_pool2)
    conv3_2 = Conv2D(256,3,activation='relu',name='conv3_2')(conv3_1)

    maxpool3 = MaxPooling2D(name='maxpool3')(conv3_2)
    conv4_1 = Conv2D(512,3,activation='relu',name='conv4_1')(maxpool3)
    conv4_2 = Conv2D(512,3,activation='relu',name='conv4_2')(conv4_1)

    maxpool4 = MaxPooling2D(name='maxpool4')(conv4_2)
    conv5_1 = Conv2D(1024,3,activation='relu',name='conv5_1')(maxpool4)
    conv5_2 = Conv2D(1024,3,activation='relu',name='conv5_2')(conv5_1)

    up5 = UpSampling2D(name='up5')(conv5_2)
    up5_conv = Conv2D(512,2,padding='same',name='up5_conv')(up5)
    conv4_feature = tf.image.resize(conv4_2,(up5_conv.shape[1],up5_conv.shape[2]),
                                    tf.image.ResizeMethod.NEAREST_NEIGHBOR,
                                    name='conv4_feature')

    concat1 = Concatenate(name='concat1')([up5_conv,conv4_feature])
    conv6_1 = Conv2D(512,3,activation='relu',name='conv6_1')(concat1)
    conv6_2 = Conv2D(512,3,activation='relu',name='conv6_2')(conv6_1)

    up6 = UpSampling2D(name='up6')(conv6_2)
    up6_conv = Conv2D(256,2,padding='same',name='up6_conv')(up6)
    conv3_feature = tf.image.resize(conv3_2,(up6_conv.shape[1],up6_conv.shape[2]),
                                    tf.image.ResizeMethod.NEAREST_NEIGHBOR,
                                    name='conv3_feature')

    concat2 = Concatenate(name='concat2')([up6_conv,conv3_feature])
    conv7_1 = Conv2D(256,3,activation='relu',name='conv7_1')(concat2)
    conv7_2 = Conv2D(256,3,activation='relu',name='conv7_2')(conv7_1)

    up7 = UpSampling2D(name='up7')(conv7_2)
    up7_conv = Conv2D(128,2,padding='same',name='up7_conv')(up7)
    conv2_feature = tf.image.resize(conv2_2,(up7_conv.shape[1],up7_conv.shape[2]),
                                    tf.image.ResizeMethod.NEAREST_NEIGHBOR,
                                    name='conv2_feature')

    concat3 = Concatenate(name='concat3')([up7_conv,conv2_feature])
    conv8_1 = Conv2D(128,3,activation='relu',name='conv8_1')(concat3)
    conv8_2 = Conv2D(128,3,activation='relu',name='conv8_2')(conv8_1)

    up8 = UpSampling2D(name='up8')(conv8_2)
    up8_conv = Conv2D(64,2,padding='same',name='up8_conv')(up8)
    conv1_feature = tf.image.resize(conv1_2,(up8_conv.shape[1],up8_conv.shape[2]),
                                    tf.image.ResizeMethod.NEAREST_NEIGHBOR,
                                    name='conv1_feature')

    concat4 = Concatenate(name='concat4')([up8_conv,conv1_feature])
    conv9_1 = Conv2D(64,3,activation='relu',name='conv9_1')(concat4)
    conv9_2 = Conv2D(64,3,activation='relu',name='conv9_2')(conv9_1)

    out = Conv2D(2,1,activation='sigmoid',name='out')(conv9_2)

    model = Model(inputs=inputs,outputs=out,name='mytf2_unet')

    return model

if __name__ == '__main__':
    model = unet()
    model.summary()

输出：

D:\program\miniconda\envs\py38_tf2\python.exe E:/paper/nets/unet.py
Model: "mytf2_unet"
__________________________________________________________________________________________________
 Layer (type)                   Output Shape         Param #     Connected to                     
==================================================================================================
 input (InputLayer)             [(None, 572, 572, 1  0           []                               
                                )]                                                                
                                                                                                  
 conv1_1 (Conv2D)               (None, 570, 570, 64  640         ['input[0][0]']                  
                                )                                                                 
                                                                                                  
 conv1_2 (Conv2D)               (None, 568, 568, 64  36928       ['conv1_1[0][0]']                
                                )                                                                 
                                                                                                  
 maxpool1 (MaxPooling2D)        (None, 284, 284, 64  0           ['conv1_2[0][0]']                
                                )                                                                 
                                                                                                  
 conv2_1 (Conv2D)               (None, 282, 282, 12  73856       ['maxpool1[0][0]']               
                                8)                                                                
                                                                                                  
 conv2_2 (Conv2D)               (None, 280, 280, 12  147584      ['conv2_1[0][0]']                
                                8)                                                                
                                                                                                  
 maxpool2 (MaxPooling2D)        (None, 140, 140, 12  0           ['conv2_2[0][0]']                
                                8)                                                                
                                                                                                  
 conv3_1 (Conv2D)               (None, 138, 138, 25  295168      ['maxpool2[0][0]']               
                                6)                                                                
                                                                                                  
 conv3_2 (Conv2D)               (None, 136, 136, 25  590080      ['conv3_1[0][0]']                
                                6)                                                                
                                                                                                  
 maxpool3 (MaxPooling2D)        (None, 68, 68, 256)  0           ['conv3_2[0][0]']                
                                                                                                  
 conv4_1 (Conv2D)               (None, 66, 66, 512)  1180160     ['maxpool3[0][0]']               
                                                                                                  
 conv4_2 (Conv2D)               (None, 64, 64, 512)  2359808     ['conv4_1[0][0]']                
                                                                                                  
 maxpool4 (MaxPooling2D)        (None, 32, 32, 512)  0           ['conv4_2[0][0]']                
                                                                                                  
 conv5_1 (Conv2D)               (None, 30, 30, 1024  4719616     ['maxpool4[0][0]']               
                                )                                                                 
                                                                                                  
 conv5_2 (Conv2D)               (None, 28, 28, 1024  9438208     ['conv5_1[0][0]']                
                                )                                                                 
                                                                                                  
 up5 (UpSampling2D)             (None, 56, 56, 1024  0           ['conv5_2[0][0]']                
                                )                                                                 
                                                                                                  
 up5_conv (Conv2D)              (None, 56, 56, 512)  2097664     ['up5[0][0]']                    
                                                                                                  
 tf.image.resize (TFOpLambda)   (None, 56, 56, 512)  0           ['conv4_2[0][0]']                
                                                                                                  
 concat1 (Concatenate)          (None, 56, 56, 1024  0           ['up5_conv[0][0]',               
                                )                                 'tf.image.resize[0][0]']        
                                                                                                  
 conv6_1 (Conv2D)               (None, 54, 54, 512)  4719104     ['concat1[0][0]']                
                                                                                                  
 conv6_2 (Conv2D)               (None, 52, 52, 512)  2359808     ['conv6_1[0][0]']                
                                                                                                  
 up6 (UpSampling2D)             (None, 104, 104, 51  0           ['conv6_2[0][0]']                
                                2)                                                                
                                                                                                  
 up6_conv (Conv2D)              (None, 104, 104, 25  524544      ['up6[0][0]']                    
                                6)                                                                
                                                                                                  
 tf.image.resize_1 (TFOpLambda)  (None, 104, 104, 25  0          ['conv3_2[0][0]']                
                                6)                                                                
                                                                                                  
 concat2 (Concatenate)          (None, 104, 104, 51  0           ['up6_conv[0][0]',               
                                2)                                'tf.image.resize_1[0][0]']      
                                                                                                  
 conv7_1 (Conv2D)               (None, 102, 102, 25  1179904     ['concat2[0][0]']                
                                6)                                                                
                                                                                                  
 conv7_2 (Conv2D)               (None, 100, 100, 25  590080      ['conv7_1[0][0]']                
                                6)                                                                
                                                                                                  
 up7 (UpSampling2D)             (None, 200, 200, 25  0           ['conv7_2[0][0]']                
                                6)                                                                
                                                                                                  
 up7_conv (Conv2D)              (None, 200, 200, 12  131200      ['up7[0][0]']                    
                                8)                                                                
                                                                                                  
 tf.image.resize_2 (TFOpLambda)  (None, 200, 200, 12  0          ['conv2_2[0][0]']                
                                8)                                                                
                                                                                                  
 concat3 (Concatenate)          (None, 200, 200, 25  0           ['up7_conv[0][0]',               
                                6)                                'tf.image.resize_2[0][0]']      
                                                                                                  
 conv8_1 (Conv2D)               (None, 198, 198, 12  295040      ['concat3[0][0]']                
                                8)                                                                
                                                                                                  
 conv8_2 (Conv2D)               (None, 196, 196, 12  147584      ['conv8_1[0][0]']                
                                8)                                                                
                                                                                                  
 up8 (UpSampling2D)             (None, 392, 392, 12  0           ['conv8_2[0][0]']                
                                8)                                                                
                                                                                                  
 up8_conv (Conv2D)              (None, 392, 392, 64  32832       ['up8[0][0]']                    
                                )                                                                 
                                                                                                  
 tf.image.resize_3 (TFOpLambda)  (None, 392, 392, 64  0          ['conv1_2[0][0]']                
                                )                                                                 
                                                                                                  
 concat4 (Concatenate)          (None, 392, 392, 12  0           ['up8_conv[0][0]',               
                                8)                                'tf.image.resize_3[0][0]']      
                                                                                                  
 conv9_1 (Conv2D)               (None, 390, 390, 64  73792       ['concat4[0][0]']                
                                )                                                                 
                                                                                                  
 conv9_2 (Conv2D)               (None, 388, 388, 64  36928       ['conv9_1[0][0]']                
                                )                                                                 
                                                                                                  
 out (Conv2D)                   (None, 388, 388, 2)  130         ['conv9_2[0][0]']                
                                                                                                  
==================================================================================================
Total params: 31,030,658
Trainable params: 31,030,658
Non-trainable params: 0
__________________________________________________________________________________________________

Process finished with exit code 0