mmagic.models.editors.nafnet.naf_avgpool2d

Module Contents

Classes

NAFAvgPool2d	Average Pooling 2D used in NAFNet.
Local_Base	Local Base class to use global average pooling.

Functions

replace_layers(model, base_size, train_size, fast_imp, ...)

Replace all layers with AvgPool2d.

class mmagic.models.editors.nafnet.naf_avgpool2d.NAFAvgPool2d(kernel_size=None, base_size=None, auto_pad=True, fast_imp=False, train_size=None)

Bases: mmengine.model.BaseModule

Average Pooling 2D used in NAFNet.

Note: this is different from the normal AvgPool2d in pytorch. According to: Improving Image Restoration by Revisiting Global Information Aggregation statistics are aggregated in a local region for each pixel rather than the global average pooling.

extra_repr() → str

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x)

mmagic.models.editors.nafnet.naf_avgpool2d.replace_layers(model, base_size, train_size, fast_imp, **kwargs)

Replace all layers with AvgPool2d.

class mmagic.models.editors.nafnet.naf_avgpool2d.Local_Base

Local Base class to use global average pooling.

Parameters

train_size – training image size

convert(*args, train_size, **kwargs)