Saving full model with pytorch ignite

username: Tanya_Boone
cooked: <p>Hello, I was able to run successfully this tutorial with pytorch ignite</p>
<p><a href="https://www.kaggle.com/hmendonca/multitask-efficientnetb4-ignite-amp-clr-aptos19" class="onebox" target="_blank" rel="nofollow noopener">https://www.kaggle.com/hmendonca/multitask-efficientnetb4-ignite-amp-clr-aptos19</a></p>
<p>but in order to do a pruning method I need to save the whole model (state dict is not useful), and I try to save with torch.save(model,‘model1.pth’) for example and I get the following message:</p>
<p>File “”, line 1, in<br>
torch.save(model,‘modelo_ejemplo1.pth’)</p>
<p>File “C:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py”, line 260, in save<br>
return _with_file_like(f, “wb”, lambda f: _save(obj, f, pickle_module, pickle_protocol))</p>
<p>File “C:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py”, line 185, in _with_file_like<br>
return body(f)</p>
<p>File “C:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py”, line 260, in<br>
return _with_file_like(f, “wb”, lambda f: _save(obj, f, pickle_module, pickle_protocol))</p>
<p>File “C:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py”, line 332, in _save<br>
pickler.dump(obj)</p>
<p>AttributeError: Can’t pickle local object ‘_initialize…patch_forward…new_fwd’</p>
<p>I’ve been looking a lot about this error message, but anyone seems to use ignite a lot…</p>
<p>I can save the model like this</p>
<p>to_save = {‘model’: model}<br>
handler = Checkpoint(to_save, DiskSaver(‘D:\Neural_Nets\EfficientNet\Pytorch\CIFAR10b’, create_dir=True,require_empty=False), n_saved=2)<br>
trainer.add_event_handler(Events.EPOCH_STARTED, handler)</p>
<p>but I can only save the weights…I want to save the whole model like I do with other codes that don’t use this pytorch ignite…can you help me with this issue? I think ignite is more efficient than my other way of training, that is why I want to change to it but if I can’t save the model there is no case <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Thanks in advance!</p>

username: vfdev-5
cooked: <p>Hi <a class="mention" href="/u/tanya_boone">@Tanya_Boone</a></p>
<blockquote>
<p>torch.save(model,‘model1.pth’)<br>
AttributeError: Can’t pickle local object ‘_initialize…patch_forward…new_fwd’</p>
</blockquote>
<p>seems like your model can not be saved with <code>torch.save</code>.</p>
<p>Maybe you need to replace some lambda function in there, if there are some…<br>
Can you share the model definition to take a look ?</p>

username: Tanya_Boone
cooked: <p>I don’t think is the model because I can save it with other type of training but when I switch to ignite is when I get this kind of error…</p>

username: sdesrozis
cooked: <p>Hi Tanya! It’s a bit far but could you share a model that you can not save with ignite ? Thank you very much for your help <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: Tanya_Boone
cooked: <p>Yes, for example, efficientnet</p>
<p>class Swish(nn.Module):<br>
def forward(self, x):<br>
return x * torch.sigmoid(x)</p>
<p>class Flatten(nn.Module):<br>
def forward(self, x):<br>
return x.reshape(x.shape[0], -1)</p>
<p>class SqueezeExcitation(nn.Module):</p>
<pre><code>def __init__(self, inplanes, se_planes):
    super(SqueezeExcitation, self).__init__()
    self.reduce_expand = nn.Sequential(
        nn.Conv2d(inplanes, se_planes, 
                  kernel_size=1, stride=1, padding=0, bias=True),
        Swish(),
        nn.Conv2d(se_planes, inplanes, 
                  kernel_size=1, stride=1, padding=0, bias=True),
        nn.Sigmoid()
    )

def forward(self, x):
    x_se = torch.mean(x, dim=(-2, -1), keepdim=True)
    x_se = self.reduce_expand(x_se)
    return x_se * x
</code></pre>
<p>class MBConv(nn.Module):<br>
def <strong>init</strong>(self, inplanes, planes, kernel_size, stride,<br>
expand_rate=1.0, se_rate=0.25,<br>
drop_connect_rate=0.2):<br>
super(MBConv, self).<strong>init</strong>()</p>
<pre><code>    expand_planes = int(inplanes * expand_rate)
    se_planes = max(1, int(inplanes * se_rate))

    self.expansion_conv = None        
    if expand_rate &gt; 1.0:
        self.expansion_conv = nn.Sequential(
            nn.Conv2d(inplanes, expand_planes, 
                      kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(expand_planes, momentum=0.01, eps=1e-3),
            Swish()
        )
        inplanes = expand_planes

    self.depthwise_conv = nn.Sequential(
        nn.Conv2d(inplanes, expand_planes,
                  kernel_size=kernel_size, stride=stride, 
                  padding=kernel_size // 2, groups=expand_planes,
                  bias=False),
        nn.BatchNorm2d(expand_planes, momentum=0.01, eps=1e-3),
        Swish()
    )

    self.squeeze_excitation = SqueezeExcitation(expand_planes, se_planes)
    
    self.project_conv = nn.Sequential(
        nn.Conv2d(expand_planes, planes, 
                  kernel_size=1, stride=1, padding=0, bias=False),
        nn.BatchNorm2d(planes, momentum=0.01, eps=1e-3),
    )

    self.with_skip = stride == 1
    self.drop_connect_rate = torch.tensor(drop_connect_rate, requires_grad=False)

def _drop_connect(self, x):        
    keep_prob = 1.0 - self.drop_connect_rate
    drop_mask = torch.rand(x.shape[0], 1, 1, 1) + keep_prob
    drop_mask = drop_mask.type_as(x)
    drop_mask.floor_()
    return drop_mask * x / keep_prob
    
def forward(self, x):
    z = x
    if self.expansion_conv is not None:
        x = self.expansion_conv(x)

    x = self.depthwise_conv(x)
    x = self.squeeze_excitation(x)
    x = self.project_conv(x)
    
    # Add identity skip
    if x.shape == z.shape and self.with_skip:            
        if self.training and self.drop_connect_rate is not None:
            self._drop_connect(x)
        x += z
    return x
</code></pre>
<p>def init_weights(module):<br>
if isinstance(module, nn.Conv2d):<br>
nn.init.kaiming_normal_(module.weight, a=0, mode=‘fan_out’)<br>
elif isinstance(module, nn.Linear):<br>
init_range = 1.0 / math.sqrt(module.weight.shape[1])<br>
nn.init.uniform_(module.weight, a=-init_range, b=init_range)</p>
<p>class EfficientNet(nn.Module):</p>
<pre><code>def _setup_repeats(self, num_repeats):
    return int(math.ceil(self.depth_coefficient * num_repeats))

def _setup_channels(self, num_channels):
    num_channels *= self.width_coefficient
    new_num_channels = math.floor(num_channels / self.divisor + 0.5) * self.divisor
    new_num_channels = max(self.divisor, new_num_channels)
    if new_num_channels &lt; 0.9 * num_channels:
        new_num_channels += self.divisor
    return new_num_channels

def __init__(self, num_classes, 
             width_coefficient=1.0,
             depth_coefficient=1.0,
             se_rate=0.25,
             dropout_rate=0.2,
             drop_connect_rate=0.2):
    super(EfficientNet, self).__init__()
    
    self.width_coefficient = width_coefficient
    self.depth_coefficient = depth_coefficient
    self.divisor = 8
            
    list_channels = [32, 16, 24, 40, 80, 112, 192, 320, 1280]
    list_channels = [self._setup_channels(c) for c in list_channels]
            
    list_num_repeats = [1, 2, 2, 3, 3, 4, 1]
    list_num_repeats = [self._setup_repeats(r) for r in list_num_repeats]        
    
    expand_rates = [1, 6, 6, 6, 6, 6, 6]
    strides = [1, 2, 2, 2, 1, 2, 1]
    kernel_sizes = [3, 3, 5, 3, 5, 5, 3]

    # Define stem:
    self.stem = nn.Sequential(
        nn.Conv2d(3, list_channels[0], kernel_size=3, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(list_channels[0], momentum=0.01, eps=1e-3),
        Swish()
    )
    
    # Define MBConv blocks
    blocks = []
    counter = 0
    num_blocks = sum(list_num_repeats)
    for idx in range(7):
        
        num_channels = list_channels[idx]
        next_num_channels = list_channels[idx + 1]
        num_repeats = list_num_repeats[idx]
        expand_rate = expand_rates[idx]
        kernel_size = kernel_sizes[idx]
        stride = strides[idx]
        drop_rate = drop_connect_rate * counter / num_blocks
        
        name = "MBConv{}_{}".format(expand_rate, counter)
        blocks.append((
            name,
            MBConv(num_channels, next_num_channels, 
                   kernel_size=kernel_size, stride=stride, expand_rate=expand_rate, 
                   se_rate=se_rate, drop_connect_rate=drop_rate)
        ))
        counter += 1
        for i in range(1, num_repeats):                
            name = "MBConv{}_{}".format(expand_rate, counter)
            drop_rate = drop_connect_rate * counter / num_blocks                
            blocks.append((
                name,
                MBConv(next_num_channels, next_num_channels, 
                       kernel_size=kernel_size, stride=1, expand_rate=expand_rate, 
                       se_rate=se_rate, drop_connect_rate=drop_rate)                                    
            ))
            counter += 1
    
    self.blocks = nn.Sequential(OrderedDict(blocks))
    
    # Define head
    self.head = nn.Sequential(
        nn.Conv2d(list_channels[-2], list_channels[-1], 
                  kernel_size=1, bias=False),
        nn.BatchNorm2d(list_channels[-1], momentum=0.01, eps=1e-3),
        Swish(),
        nn.AdaptiveAvgPool2d(1),
        Flatten(),
        nn.Dropout(p=dropout_rate),
        nn.Linear(list_channels[-1], num_classes)
    )

    self.apply(init_weights)
    
def forward(self, x):
    f = self.stem(x)
    f = self.blocks(f)
    y = self.head(f)
    return y
</code></pre>
<p>model = EfficientNet(num_classes=10,<br>
width_coefficient=1.4, depth_coefficient=1.8,<br>
dropout_rate=0.4)<br>
resolution = 380<br>
img_stats  = [[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]]</p>
<p>but I managed to save it with this library called dill</p>
<p>For example,</p>
<p>import dill</p>
<h1>Dill routine</h1>
<p>model_copy=dill.dumps(model)<br>
torch.save(model_copy,‘model_ignite_original.pt’)</p>
<p>and I dont get any errors…</p>
<p>when I want to load it I do it like this</p>
<pre><code>model1 = torch.load(model_name)
model=dill.loads(model1)</code></pre>

username: sdesrozis
cooked: <p>Sorry for this late response. I will look more precisely your code, thank very much !</p>