UNet with VGG16 and batch normalization

I’m getting very low accuracy for semantic segmentation problem using Unet+resnet34 (0,84 tops).

Having tried a model on supervise.ly based on Unet with VGG16 and batch normalization, with 0,95 accuracy, I tried to reproduce it with fastai.

Can someone understand why I’m getting this error?

The code runs ok changing the model to resnet34:

learn = unet_learner(data, models.vgg16_bn, metrics=metrics, wd=wd)
learn.fit_one_cycle(5, slice(1e-5))

 Interrupted
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-49-be1ab4476b35> in <module>
----> 1 learn.fit_one_cycle(5, slice(lr))

~/anaconda3/envs/fastai/lib/python3.7/site-packages/fastai/train.py in fit_one_cycle(learn, cyc_len, max_lr, moms, div_factor, pct_start, final_div, wd, callbacks, tot_epochs, start_epoch)
     20     callbacks.append(OneCycleScheduler(learn, max_lr, moms=moms, div_factor=div_factor, pct_start=pct_start,
     21                                        final_div=final_div, tot_epochs=tot_epochs, start_epoch=start_epoch))
---> 22     learn.fit(cyc_len, max_lr, wd=wd, callbacks=callbacks)
     23 
     24 def lr_find(learn:Learner, start_lr:Floats=1e-7, end_lr:Floats=10, num_it:int=100, stop_div:bool=True, wd:float=None):

~/anaconda3/envs/fastai/lib/python3.7/site-packages/fastai/basic_train.py in fit(self, epochs, lr, wd, callbacks)
    194         callbacks = [cb(self) for cb in self.callback_fns] + listify(callbacks)
    195         if defaults.extra_callbacks is not None: callbacks += defaults.extra_callbacks
--> 196         fit(epochs, self, metrics=self.metrics, callbacks=self.callbacks+callbacks)
    197 
    198     def create_opt(self, lr:Floats, wd:Floats=0.)->None:

~/anaconda3/envs/fastai/lib/python3.7/site-packages/fastai/basic_train.py in fit(epochs, learn, callbacks, metrics)
     98             for xb,yb in progress_bar(learn.data.train_dl, parent=pbar):
     99                 xb, yb = cb_handler.on_batch_begin(xb, yb)
--> 100                 loss = loss_batch(learn.model, xb, yb, learn.loss_func, learn.opt, cb_handler)
    101                 if cb_handler.on_batch_end(loss): break
    102 

~/anaconda3/envs/fastai/lib/python3.7/site-packages/fastai/basic_train.py in loss_batch(model, xb, yb, loss_func, opt, cb_handler)
     23     if not is_listy(xb): xb = [xb]
     24     if not is_listy(yb): yb = [yb]
---> 25     out = model(*xb)
     26     out = cb_handler.on_loss_begin(out)
     27 

~/anaconda3/envs/fastai/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

~/anaconda3/envs/fastai/lib/python3.7/site-packages/fastai/layers.py in forward(self, x)
    153         for l in self.layers:
    154             res.orig = x
--> 155             nres = l(res)
    156             # We have to remove res.orig to avoid hanging refs and therefore memory leaks
    157             res.orig = None

~/anaconda3/envs/fastai/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

~/anaconda3/envs/fastai/lib/python3.7/site-packages/fastai/layers.py in forward(self, x)
    170         self.dense=dense
    171 
--> 172     def forward(self, x): return torch.cat([x,x.orig], dim=1) if self.dense else (x+x.orig)
    173 
    174 def res_block(nf, dense:bool=False, norm_type:Optional[NormType]=NormType.Batch, bottle:bool=False, **conv_kwargs):

RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 1. Got 258 and 16 in dimension 2 at /opt/conda/conda-bld/pytorch_1549636813070/work/aten/src/THC/generic/THCTensorMath.cu:83

Right now only ResNets can be used (I had the same problem too :wink: ):

Since fastai supports resnets only you can find VGG16 UNet architecture implementation here. You can create your own model and pass it through the Learner.

Is there a way I can propose to include information about this limitation in the https://docs.fast.ai/vision.learner.html#unet_learner?