基于EfficientNetB4的皮肤病分类模型实现疑问与OOM问题求助

1. argmax操作的实现方式 #

根据论文逻辑，argmax是推理阶段的操作，并非训练时的网络层：

• 训练阶段：将7个辅助分类器的输出（已通过softmax得到概率分布）与最终分类器的输出求和，得到综合概率分布，以此计算损失完成训练。
• 推理阶段：对求和后的综合概率分布执行argmax操作，选取概率最大的类别作为最终预测结果。
  不能将argmax加入网络结构作为训练层，因为argmax是不可导操作，会阻断梯度反向传播，导致模型无法完成训练。

2. ResourceExhaustedError(OOM)内存耗尽错误分析与解决 #

错误原因 #

• EfficientNet-B4本身参数量、计算量较大，加上7个包含Dense(1024)的辅助分类器，大幅增加了显存占用。
• 输入尺寸380x380+批量大小30，GPU显存无法容纳模型参数、中间特征图及训练时的梯度数据。错误日志中提到的shape[32,190,190,144]特征图，正是显存占用过高的直接体现。

解决方法 #

• 降低批量大小：将batch_size从30调整为8、16等更小的值，根据GPU显存容量适配。
• 缩小输入尺寸：改用EfficientNet默认输入尺寸224x224，或299x299，减少中间特征图的显存占用。
• 冻结预训练层：先冻结base_model的大部分层，仅训练辅助分类器和顶部全连接层，待模型收敛后再解冻微调，减少训练时的梯度计算内存开销。
• 简化辅助分类器：将Dense(1024)改为更小的维度（如256、512），降低参数量与显存占用。
• 启用混合精度训练：添加以下代码开启混合精度，减少显存占用：

  import tensorflow as tf
  tf.keras.mixed_precision.set_global_policy('mixed_float16')
  

• 清理显存：训练前执行tf.keras.backend.clear_session()，释放之前模型占用的显存。

你的模型实现代码 #

classes = 7
base_model = keras.applications.EfficientNetB4(weights="imagenet",include_top=False, input_shape=(380,380,3))
avg = keras.layers.GlobalAveragePooling2D()(base_model.output)
cal_layer = keras.layers.Dense(classes, activation="softmax")(avg)

def auxilary_classifier(input, n):
  x = keras.layers.AveragePooling2D(pool_size=(5,5), strides=3)(input)
  x = keras.layers.Conv2D(128, kernel_size=1, padding='same', activation='relu')(x)
  x = keras.layers.Flatten()(x)
  x = keras.layers.BatchNormalization()(x)
  x = keras.layers.Dense(1024, activation='relu')(x)
  x = keras.layers.Dense(n, activation='softmax')(x)
  return x


# auxiliary classifier 1
block1 = base_model.get_layer('block1b_add').output
aux_head_1 = auxilary_classifier(block1, classes)

# auxiliary classifier 2
block2 = base_model.get_layer('block2d_add').output
aux_head_2 = auxilary_classifier(block2, classes)

# auxiliary classifier 3
block3 = base_model.get_layer('block3d_add').output
aux_head_3 = auxilary_classifier(block3, classes)

# auxiliary classifier 4
block4 = base_model.get_layer('block4f_add').output
aux_head_4 = auxilary_classifier(block4, classes)

# auxiliary classifier 5
block5 = base_model.get_layer('block5f_add').output
aux_head_5 = auxilary_classifier(block5, classes)

# auxiliary classifier 6
block6 = base_model.get_layer('block6h_add').output
aux_head_6 = auxilary_classifier(block6, classes)

# auxiliary classifier 7
block7 = base_model.get_layer('block7b_add').output
aux_head_7 = auxilary_classifier(block7, classes)

sum_final = keras.layers.Add()([cal_layer,aux_head_1,aux_head_2,aux_head_3,aux_head_4,aux_head_5,aux_head_6,aux_head_7])

model = keras.Model(base_model.input, sum_final)


optimizer = keras.optimizers.Nadam(learning_rate=0.01)
model.compile(loss="sparse_categorical_crossentropy",optimizer=optimizer,metrics=["accuracy"])
history = model.fit(train, validation_data=valid, epochs=30, verbose=2, batch_size=30)

错误日志 #

Epoch 1/30

---------------------------------------------------------------------------

ResourceExhaustedError                    Traceback (most recent call last)

<ipython-input-9-91732b783ef6> in <cell line: 6>()
      4 optimizer = keras.optimizers.Nadam(learning_rate=0.01)
      5 model.compile(loss="sparse_categorical_crossentropy",optimizer=optimizer,metrics=["accuracy"])
----> 6 history = model.fit(train, validation_data=valid, epochs=30, verbose=2, batch_size=30, callbacks=[reduce_lr])

1 frames

/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
     68             # To get the full stack trace, call:
     69             # `tf.debugging.disable_traceback_filtering()`
---> 70             raise e.with_traceback(filtered_tb) from None
     71         finally:
     72             del filtered_tb

/usr/local/lib/python3.9/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     50   try:
     51     ctx.ensure_initialized()
---> 52     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     53                                         inputs, attrs, num_outputs)
     54   except core._NotOkStatusException as e:

ResourceExhaustedError: Graph execution error:

Detected at node 'model_1/block2a_expand_conv/Conv2D' defined at (most recent call last):
    File "/usr/lib/python3.9/runpy.py", line 197, in _run_module_as_main
      return _run_code(code, main_globals, None,
    File "/usr/lib/python3.9/runpy.py", line 87, in _run_code
      exec(code, run_globals)
    File "/usr/local/lib/python3.9/dist-packages/ipykernel_launcher.py", line 16, in <module>
      app.launch_new_instance()
    File "/usr/local/lib/python3.9/dist-packages/traitlets/config/application.py", line 992, in launch_instance
      app.start()
    File "/usr/local/lib/python3.9/dist-packages/ipykernel/kernelapp.py", line 619, in start
      self.io_loop.start()
    File "/usr/local/lib/python3.9/dist-packages/tornado/platform/asyncio.py", line 215, in start
      self.asyncio_loop.run_forever()
    File "/usr/lib/python3.9/asyncio/base_events.py", line 601, in run_forever
      self._run_once()
    File "/usr/lib/python3.9/asyncio/base_events.py", line 1905, in _run_once
      handle._run()
    File "/usr/lib/python3.9/asyncio/events.py", line 80, in _run
      self._context.run(self._callback, *self._args)
    File "/usr/local/lib/python3.9/dist-packages/tornado/ioloop.py", line 687, in <lambda>
      lambda f: self._run_callback(functools.partial(callback, future))
    File "/usr/local/lib/python3.9/dist-packages/tornado/ioloop.py", line 740, in _run_callback
      ret = callback()
    File "/usr/local/lib/python3.9/dist-packages/tornado/gen.py", line 821, in inner
      self.ctx_run(self.run)
    File "/usr/local/lib/python3.9/dist-packages/tornado/gen.py", line 782, in run
      yielded = self.gen.send(value)
    File "/usr/local/lib/python3.9/dist-packages/ipykernel/kernelbase.py", line 361, in process_one
      yield gen.maybe_future(dispatch(*args))
    File "/usr/local/lib/python3.9/dist-packages/tornado/gen.py", line 234, in wrapper
      yielded = ctx_run(next, result)
    File "/usr/local/lib/python3.9/dist-packages/ipykernel/kernelbase.py", line 261, in dispatch_shell
      yield gen.maybe_future(handler(stream, idents, msg))
    File "/usr/local/lib/python3.9/dist-packages/tornado/gen.py", line 234, in wrapper
      yielded = ctx_run(next, result)
    File "/usr/local/lib/python3.9/dist-packages/ipykernel/kernelbase.py", line 539, in execute_request
      self.do_execute(
    File "/usr/local/lib/python3.9/dist-packages/tornado/gen.py", line 234, in wrapper
      yielded = ctx_run(next, result)
    File "/usr/local/lib/python3.9/dist-packages/ipykernel/ipkernel.py", line 302, in do_execute
      res = shell.run_cell(code, store_history=store_history, silent=silent)
    File "/usr/local/lib/python3.9/dist-packages/ipykernel/zmqshell.py", line 539, in run_cell
      return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/IPython/core/interactiveshell.py", line 2975, in run_cell
      result = self._run_cell(
    File "/usr/local/lib/python3.9/dist-packages/IPython/core/interactiveshell.py", line 3030, in _run_cell
      return runner(coro)
    File "/usr/local/lib/python3.9/dist-packages/IPython/core/async_helpers.py", line 78, in _pseudo_sync_runner
      coro.send(None)
    File "/usr/local/lib/python3.9/dist-packages/IPython/core/interactiveshell.py", line 3257, in run_cell_async
      has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
    File "/usr/local/lib/python3.9/dist-packages/IPython/core/interactiveshell.py", line 3473, in run_ast_nodes
      if (await self.run_code(code, result,  async_=asy)):
    File "/usr/local/lib/python3.9/dist-packages/IPython/core/interactiveshell.py", line 3553, in run_code
      exec(code_obj, self.user_global_ns, self.user_ns)
    File "<ipython-input-9-91732b783ef6>", line 6, in <cell line: 6>
      history = model.fit(train, validation_data=valid, epochs=30, verbose=2, batch_size=30, callbacks=[reduce_lr])
    File "/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/keras/engine/training.py", line 1685, in fit
      tmp_logs = self.train_function(iterator)
    File "/usr/local/lib/python3.9/dist-packages/keras/engine/training.py", line 1284, in train_function
      return step_function(self, iterator)
    File "/usr/local/lib/python3.9/dist-packages/keras/engine/training.py", line 1268, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.9/dist-packages/keras/engine/training.py", line 1249, in run_step
      outputs = model.train_step(data)
    File "/usr/local/lib/python3.9/dist-packages/keras/engine/training.py", line 1050, in train_step
      y_pred = self(x, training=True)
    File "/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/keras/engine/training.py", line 558, in __call__
      return super().__call__(*args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/keras/engine/base_layer.py", line 1145, in __call__
      outputs = call_fn(inputs, *args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py", line 96, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/keras/engine/functional.py", line 512, in call
      return self._run_internal_graph(inputs, training=training, mask=mask)
    File "/usr/local/lib/python3.9/dist-packages/keras/engine/functional.py", line 669, in _run_internal_graph
      outputs = node.layer(*args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/keras/engine/base_layer.py", line 1145, in __call__
      outputs = call_fn(inputs, *args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py", line 96, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.9/dist-packages/keras/layers/convolutional/base_conv.py", line 290, in call
      outputs = self.convolution_op(inputs, self.kernel)
    File "/usr/local/lib/python3.9/dist-packages/keras/layers/convolutional/base_conv.py", line 262, in convolution_op
      return tf.nn.convolution(
Node: 'model_1/block2a_expand_conv/Conv2D'
OOM when allocating tensor with shape[32,190,190,144] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
     [[{{node model_1/block2a_expand_conv/Conv2D}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info. This isn't available when running in Eager mode.
 [Op:__inference_train_function_115485]

内容的提问来源于stack exchange，提问作者Martin