如何减少我的tensorflow模型所用的GPU？
我现在正在使用Kaggle的笔记本环境，想要构建一个用于图像分割的U-Net模型。当我构建这个模型时，显示模型的大小是2.17MB，但它却占用了15.4GB的GPU内存。这是为什么呢？有没有办法减少这个内存占用？使用pytorch构建模型会有所帮助吗？我对这些还不太熟悉，所以任何帮助都非常感谢！
这是我的模型代码，方便大家复现。
def build_ARes_UNet(IMG_HEIGHT, IMG_WIDHT, IMG_CHANNELS, FILTER_NUM, FILTER_SIZE, DROPOUT, ACTIVATION, RESIDUAL, ATTENTION):
inputs = tf.keras.layers.Input((IMG_WIDHT, IMG_HEIGHT, IMG_CHANNELS))
# s = tf.keras.layers.Lambda(lambda x:x/255)(inputs)
#layer 1
c1 = res_conv_block(x=inputs, size=FILTER_NUM, f_size=FILTER_SIZE, dropout=DROPOUT, activation=ACTIVATION, res=RESIDUAL) #16 FOR FITLER_NUM=16
p1 = MaxPooling2D((2,2))(c1)
#layer 2
c2 = res_conv_block(x=p1, size=FILTER_NUM*2, f_size=FILTER_SIZE, dropout=DROPOUT, activation=ACTIVATION, res=RESIDUAL) #32
p2 = MaxPooling2D((2,2))(c2)
#layer 3
c3 = res_conv_block(x=p2, size=FILTER_NUM*4, f_size=FILTER_SIZE, dropout=DROPOUT, activation=ACTIVATION, res=RESIDUAL) #64
p3 = MaxPooling2D((2,2))(c3)
#layer 4
c4 = res_conv_block(x=p3, size=FILTER_NUM*8, f_size=FILTER_SIZE, dropout=DROPOUT, activation=ACTIVATION, res=RESIDUAL) #128
p4 = MaxPooling2D(pool_size=(2, 2))(c4)
#layer 5
c5 = Conv2D(FILTER_NUM*16, (3, 3), activation='relu', padding='same')(p4) #256
c5 = Dropout(0.3)(c5)
c5 = Conv2D(FILTER_NUM*16, (3, 3), activation='relu', padding='same')(c5)
#Expansive path
#layer 6
u6 = UpSampling2D(size=(2,2), data_format="channels_last")(c5)
if ATTENTION[0]:
gate6 = gate_block(c5, FILTER_NUM*8) #128
attention6 = attention_block(c4, gate6, FILTER_NUM*8)
# u6 = Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(c5)
u6 = concatenate([u6, attention6], axis=3)
else:
u6 = concatenate([u6,c4])
c6 = res_conv_block(x=u6, size=FILTER_NUM*8, f_size=FILTER_SIZE, dropout=DROPOUT, activation=ACTIVATION, res=RESIDUAL) #128
#layer 7
u7 = UpSampling2D(size=(2,2), data_format="channels_last")(c6)
if ATTENTION[1]:
gate7 = gate_block(c6, FILTER_NUM*4)
attention7 = attention_block(c3, gate7, FILTER_NUM*4)
# u7 = Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(c6)
u7 = concatenate([u7, attention7])
else:
u7 = concatenate([u7,c3])
c7 = res_conv_block(x=u7, size=FILTER_NUM*4, f_size=FILTER_SIZE, dropout=DROPOUT, activation=ACTIVATION, res=RESIDUAL) #64
#layer 8
u8 = UpSampling2D(size=(2,2), data_format="channels_last")(c7)
if ATTENTION[2]:
gate8 = gate_block(c7, FILTER_NUM*2)
attention8 = attention_block(c2, gate8, FILTER_NUM*2)
# u8 = Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(c7)
u8 = concatenate([u8, attention8])
else:
u8 = concatenate([u8,c2])
c8 = res_conv_block(x=u8, size=FILTER_NUM*2, f_size=FILTER_SIZE, dropout=DROPOUT, activation=ACTIVATION, res=RESIDUAL) #32
#layer 9
u9 = UpSampling2D(size=(2,2), data_format="channels_last")(c8)
if ATTENTION[3]:
gate9 = gate_block(c8, FILTER_NUM)
attention9 = attention_block(c1, gate9, FILTER_NUM)
# u9 = Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same')(c8)
u9 = concatenate([u9, attention9], axis=3)
else:
u9 = concatenate([u9,c1])
c9 = res_conv_block(x=u9, size=FILTER_NUM, f_size=FILTER_SIZE, dropout=DROPOUT, activation=ACTIVATION, res=RESIDUAL) #16
outputs = Conv2D(1, (1, 1), activation='sigmoid', padding = 'same')(c9)
model = tf.keras.models.Model(inputs = inputs, outputs=outputs)
return model
res_conv_block的定义如下：
def res_conv_block(x, size, f_size, dropout, activation, res):
c1 = Conv2D(size,(f_size, f_size), padding='same',activation=activation)(x)
c1 = Dropout(dropout)(c1)
c1 = Conv2D(size,(f_size, f_size), padding='same', activation=activation)(c1)
if not res:
return c1
else:
shortcut = Conv2D(size, kernel_size = (1,1), padding ='same')(x)
respath = tf.keras.layers.Add()([shortcut, c1])
respath = Activation(activation)(c1)
return respath
这是我模型总结的输出：
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type) ┃ Output Shape ┃ Param # ┃ Connected to ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩
│ input_layer │ (None, 256, 256, │ 0 │ - │
│ (InputLayer) │ 3) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d (Conv2D) │ (None, 256, 256, │ 224 │ input_layer[0][0] │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout (Dropout) │ (None, 256, 256, │ 0 │ conv2d[0][0] │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_1 (Conv2D) │ (None, 256, 256, │ 584 │ dropout[0][0] │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation │ (None, 256, 256, │ 0 │ conv2d_1[0][0] │
│ (Activation) │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ max_pooling2d │ (None, 128, 128, │ 0 │ activation[0][0] │
│ (MaxPooling2D) │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_3 (Conv2D) │ (None, 128, 128, │ 1,168 │ max_pooling2d[0]… │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout_1 (Dropout) │ (None, 128, 128, │ 0 │ conv2d_3[0][0] │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_4 (Conv2D) │ (None, 128, 128, │ 2,320 │ dropout_1[0][0] │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_1 │ (None, 128, 128, │ 0 │ conv2d_4[0][0] │
│ (Activation) │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ max_pooling2d_1 │ (None, 64, 64, │ 0 │ activation_1[0][… │
│ (MaxPooling2D) │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_6 (Conv2D) │ (None, 64, 64, │ 4,640 │ max_pooling2d_1[… │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout_2 (Dropout) │ (None, 64, 64, │ 0 │ conv2d_6[0][0] │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_7 (Conv2D) │ (None, 64, 64, │ 9,248 │ dropout_2[0][0] │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_2 │ (None, 64, 64, │ 0 │ conv2d_7[0][0] │
│ (Activation) │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ max_pooling2d_2 │ (None, 32, 32, │ 0 │ activation_2[0][… │
│ (MaxPooling2D) │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_9 (Conv2D) │ (None, 32, 32, │ 18,496 │ max_pooling2d_2[… │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout_3 (Dropout) │ (None, 32, 32, │ 0 │ conv2d_9[0][0] │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_10 (Conv2D) │ (None, 32, 32, │ 36,928 │ dropout_3[0][0] │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_3 │ (None, 32, 32, │ 0 │ conv2d_10[0][0] │
│ (Activation) │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ max_pooling2d_3 │ (None, 16, 16, │ 0 │ activation_3[0][… │
│ (MaxPooling2D) │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_12 (Conv2D) │ (None, 16, 16, │ 73,856 │ max_pooling2d_3[… │
│ │ 128) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout_4 (Dropout) │ (None, 16, 16, │ 0 │ conv2d_12[0][0] │
│ │ 128) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_13 (Conv2D) │ (None, 16, 16, │ 147,584 │ dropout_4[0][0] │
│ │ 128) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_14 (Conv2D) │ (None, 16, 16, │ 8,256 │ conv2d_13[0][0] │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_16 (Conv2D) │ (None, 16, 16, │ 4,160 │ conv2d_14[0][0] │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_transpose │ (None, 16, 16, │ 36,928 │ conv2d_16[0][0] │
│ (Conv2DTranspose) │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_15 (Conv2D) │ (None, 16, 16, │ 4,160 │ activation_3[0][… │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ add_4 (Add) │ (None, 16, 16, │ 0 │ conv2d_transpose… │
│ │ 64) │ │ conv2d_15[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_4 │ (None, 16, 16, │ 0 │ add_4[0][0] │
│ (Activation) │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_17 (Conv2D) │ (None, 16, 16, 1) │ 65 │ activation_4[0][… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_5 │ (None, 16, 16, 1) │ 0 │ conv2d_17[0][0] │
│ (Activation) │ │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ up_sampling2d_1 │ (None, 32, 32, 1) │ 0 │ activation_5[0][… │
│ (UpSampling2D) │ │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ lambda (Lambda) │ (None, 32, 32, │ 0 │ up_sampling2d_1[… │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ multiply (Multiply) │ (None, 32, 32, │ 0 │ lambda[0][0], │
│ │ 64) │ │ activation_3[0][… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_18 (Conv2D) │ (None, 32, 32, │ 4,160 │ multiply[0][0] │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ up_sampling2d │ (None, 32, 32, │ 0 │ conv2d_13[0][0] │
│ (UpSampling2D) │ 128) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ batch_normalization │ (None, 32, 32, │ 256 │ conv2d_18[0][0] │
│ (BatchNormalizatio… │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ concatenate │ (None, 32, 32, │ 0 │ up_sampling2d[0]… │
│ (Concatenate) │ 192) │ │ batch_normalizat… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_19 (Conv2D) │ (None, 32, 32, │ 110,656 │ concatenate[0][0] │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout_5 (Dropout) │ (None, 32, 32, │ 0 │ conv2d_19[0][0] │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_20 (Conv2D) │ (None, 32, 32, │ 36,928 │ dropout_5[0][0] │
│ │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_6 │ (None, 32, 32, │ 0 │ conv2d_20[0][0] │
│ (Activation) │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_22 (Conv2D) │ (None, 32, 32, │ 2,080 │ activation_6[0][… │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_24 (Conv2D) │ (None, 32, 32, │ 1,056 │ conv2d_22[0][0] │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_transpose_1 │ (None, 32, 32, │ 9,248 │ conv2d_24[0][0] │
│ (Conv2DTranspose) │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_23 (Conv2D) │ (None, 32, 32, │ 1,056 │ activation_2[0][… │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ add_6 (Add) │ (None, 32, 32, │ 0 │ conv2d_transpose… │
│ │ 32) │ │ conv2d_23[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_7 │ (None, 32, 32, │ 0 │ add_6[0][0] │
│ (Activation) │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_25 (Conv2D) │ (None, 32, 32, 1) │ 33 │ activation_7[0][… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_8 │ (None, 32, 32, 1) │ 0 │ conv2d_25[0][0] │
│ (Activation) │ │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ up_sampling2d_3 │ (None, 64, 64, 1) │ 0 │ activation_8[0][… │
│ (UpSampling2D) │ │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ lambda_1 (Lambda) │ (None, 64, 64, │ 0 │ up_sampling2d_3[… │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ multiply_1 │ (None, 64, 64, │ 0 │ lambda_1[0][0], │
│ (Multiply) │ 32) │ │ activation_2[0][… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_26 (Conv2D) │ (None, 64, 64, │ 1,056 │ multiply_1[0][0] │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ up_sampling2d_2 │ (None, 64, 64, │ 0 │ activation_6[0][… │
│ (UpSampling2D) │ 64) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ batch_normalizatio… │ (None, 64, 64, │ 128 │ conv2d_26[0][0] │
│ (BatchNormalizatio… │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ concatenate_1 │ (None, 64, 64, │ 0 │ up_sampling2d_2[… │
│ (Concatenate) │ 96) │ │ batch_normalizat… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_27 (Conv2D) │ (None, 64, 64, │ 27,680 │ concatenate_1[0]… │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout_6 (Dropout) │ (None, 64, 64, │ 0 │ conv2d_27[0][0] │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_28 (Conv2D) │ (None, 64, 64, │ 9,248 │ dropout_6[0][0] │
│ │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_9 │ (None, 64, 64, │ 0 │ conv2d_28[0][0] │
│ (Activation) │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_30 (Conv2D) │ (None, 64, 64, │ 528 │ activation_9[0][… │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_32 (Conv2D) │ (None, 64, 64, │ 272 │ conv2d_30[0][0] │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_transpose_2 │ (None, 64, 64, │ 2,320 │ conv2d_32[0][0] │
│ (Conv2DTranspose) │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_31 (Conv2D) │ (None, 64, 64, │ 272 │ activation_1[0][… │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ add_8 (Add) │ (None, 64, 64, │ 0 │ conv2d_transpose… │
│ │ 16) │ │ conv2d_31[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_10 │ (None, 64, 64, │ 0 │ add_8[0][0] │
│ (Activation) │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_33 (Conv2D) │ (None, 64, 64, 1) │ 17 │ activation_10[0]… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_11 │ (None, 64, 64, 1) │ 0 │ conv2d_33[0][0] │
│ (Activation) │ │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ up_sampling2d_5 │ (None, 128, 128, │ 0 │ activation_11[0]… │
│ (UpSampling2D) │ 1) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ lambda_2 (Lambda) │ (None, 128, 128, │ 0 │ up_sampling2d_5[… │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ multiply_2 │ (None, 128, 128, │ 0 │ lambda_2[0][0], │
│ (Multiply) │ 16) │ │ activation_1[0][… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_34 (Conv2D) │ (None, 128, 128, │ 272 │ multiply_2[0][0] │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ up_sampling2d_4 │ (None, 128, 128, │ 0 │ activation_9[0][… │
│ (UpSampling2D) │ 32) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ batch_normalizatio… │ (None, 128, 128, │ 64 │ conv2d_34[0][0] │
│ (BatchNormalizatio… │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ concatenate_2 │ (None, 128, 128, │ 0 │ up_sampling2d_4[… │
│ (Concatenate) │ 48) │ │ batch_normalizat… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_35 (Conv2D) │ (None, 128, 128, │ 6,928 │ concatenate_2[0]… │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout_7 (Dropout) │ (None, 128, 128, │ 0 │ conv2d_35[0][0] │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_36 (Conv2D) │ (None, 128, 128, │ 2,320 │ dropout_7[0][0] │
│ │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_12 │ (None, 128, 128, │ 0 │ conv2d_36[0][0] │
│ (Activation) │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_38 (Conv2D) │ (None, 128, 128, │ 136 │ activation_12[0]… │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_40 (Conv2D) │ (None, 128, 128, │ 72 │ conv2d_38[0][0] │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_transpose_3 │ (None, 128, 128, │ 584 │ conv2d_40[0][0] │
│ (Conv2DTranspose) │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_39 (Conv2D) │ (None, 128, 128, │ 72 │ activation[0][0] │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ add_10 (Add) │ (None, 128, 128, │ 0 │ conv2d_transpose… │
│ │ 8) │ │ conv2d_39[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_13 │ (None, 128, 128, │ 0 │ add_10[0][0] │
│ (Activation) │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_41 (Conv2D) │ (None, 128, 128, │ 9 │ activation_13[0]… │
│ │ 1) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_14 │ (None, 128, 128, │ 0 │ conv2d_41[0][0] │
│ (Activation) │ 1) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ up_sampling2d_7 │ (None, 256, 256, │ 0 │ activation_14[0]… │
│ (UpSampling2D) │ 1) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ lambda_3 (Lambda) │ (None, 256, 256, │ 0 │ up_sampling2d_7[… │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ multiply_3 │ (None, 256, 256, │ 0 │ lambda_3[0][0], │
│ (Multiply) │ 8) │ │ activation[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_42 (Conv2D) │ (None, 256, 256, │ 72 │ multiply_3[0][0] │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ up_sampling2d_6 │ (None, 256, 256, │ 0 │ activation_12[0]… │
│ (UpSampling2D) │ 16) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ batch_normalizatio… │ (None, 256, 256, │ 32 │ conv2d_42[0][0] │
│ (BatchNormalizatio… │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ concatenate_3 │ (None, 256, 256, │ 0 │ up_sampling2d_6[… │
│ (Concatenate) │ 24) │ │ batch_normalizat… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_43 (Conv2D) │ (None, 256, 256, │ 1,736 │ concatenate_3[0]… │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dropout_8 (Dropout) │ (None, 256, 256, │ 0 │ conv2d_43[0][0] │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_44 (Conv2D) │ (None, 256, 256, │ 584 │ dropout_8[0][0] │
│ │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ activation_15 │ (None, 256, 256, │ 0 │ conv2d_44[0][0] │
│ (Activation) │ 8) │ │ │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ conv2d_46 (Conv2D) │ (None, 256, 256, │ 9 │ activation_15[0]… │
│ │ 1) │ │ │
└─────────────────────┴───────────────────┴────────────┴───────────────────┘
Total params: 568,501 (2.17 MB)
Trainable params: 568,261 (2.17 MB)
Non-trainable params: 240 (960.00 B)
这是Kaggle显示的我的内存使用情况：
图片链接
我的图像尺寸是256x256x3，我使用的是学习率为0.001的Adam优化器，具体如下：
model_ATEST = build_ARes_UNet(IMG_HEIGHT, IMG_WIDHT, IMG_CHANNELS, FILTER_NUM=8, FILTER_SIZE=3, DROPOUT=0.2, ACTIVATION='relu', RESIDUAL=1, ATTENTION=[1,1,1,1])
model_ATEST.compile(optimizer=tf.keras.optimizers.Adagrad(0.001), loss=dice_loss, metrics=[dice_coef])
model_ATEST.summary()
1 个回答
0
默认情况下，TensorFlow 会尽可能多地使用 GPU 的内存，这样可以让它运行得更快，而不是一点一点地分配内存。如果你不和别人共享 GPU，或者不需要它来做其他任务，这样做是没问题的。
不过，你可以选择关闭这种行为，让 TensorFlow 只分配它真正需要的内存。你可以看看 这个例子。其中一个重要的代码行是 这一行。根据文档的说明：
physical_devices = tf.config.list_physical_devices('GPU') # get a list of gpus
try:
tf.config.experimental.set_memory_growth(physical_devices[0], True) # set the growth for the first one
except:
# Invalid device or cannot modify virtual devices once initialized.
pass
如果你只有一块 GPU，这样设置就够了。如果有多块 GPU，你可以遍历 physical_devices 列表，为每一块 GPU 设置内存增长。