Honestly I don’t understand where the issue is coming from either, this was just a quick fix. The error seems to be coming from fastai side though, According to sgugger on fastai’s forum :
" Ah! I think this might be due to our `data_collate` default function, which collected the `data` inside your tensor instead of just grabbing your tensor.
Why that didn’t release memory is beyond me, but I think if you pass to the call to `DataBunch` the regular pytorch collate function (which is `torch.utils.data.dataloader.default_collate` ) you won’t have a memory leak."
The Image source code is the following :
class Image(ItemBase):
"Support applying transforms to image data in `px`."
def __init__(self, px:Tensor):
self._px = px
self._logit_px=None
self._flow=None
self._affine_mat=None
self.sample_kwargs = {}
def set_sample(self, **kwargs)->'ImageBase':
"Set parameters that control how we `grid_sample` the image after transforms are applied."
self.sample_kwargs = kwargs
return self
def clone(self):
"Mimic the behavior of torch.clone for `Image` objects."
return self.__class__(self.px.clone())
@property
def shape(self)->Tuple[int,int,int]: return self._px.shape
@property
def size(self)->Tuple[int,int]: return self.shape[-2:]
@property
def device(self)->torch.device: return self._px.device
def __repr__(self): return f'{self.__class__.__name__} {tuple(self.shape)}'
def _repr_png_(self): return self._repr_image_format('png')
def _repr_jpeg_(self): return self._repr_image_format('jpeg')
def _repr_image_format(self, format_str):
with BytesIO() as str_buffer:
plt.imsave(str_buffer, image2np(self.px), format=format_str)
return str_buffer.getvalue()
def apply_tfms(self, tfms:TfmList, do_resolve:bool=True, xtra:Optional[Dict[Callable,dict]]=None,
size:Optional[Union[int,TensorImageSize]]=None, resize_method:ResizeMethod=None,
mult:int=None, padding_mode:str='reflection', mode:str='bilinear', remove_out:bool=True)->TensorImage:
"Apply all `tfms` to the `Image`, if `do_resolve` picks value for random args."
if not (tfms or xtra or size): return self
tfms = listify(tfms)
xtra = ifnone(xtra, {})
default_rsz = ResizeMethod.SQUISH if (size is not None and is_listy(size)) else ResizeMethod.CROP
resize_method = ifnone(resize_method, default_rsz)
if resize_method <= 2 and size is not None: tfms = self._maybe_add_crop_pad(tfms)
tfms = sorted(tfms, key=lambda o: o.tfm.order)
if do_resolve: _resolve_tfms(tfms)
x = self.clone()
x.set_sample(padding_mode=padding_mode, mode=mode, remove_out=remove_out)
if size is not None:
crop_target = _get_crop_target(size, mult=mult)
if resize_method in (ResizeMethod.CROP,ResizeMethod.PAD):
target = _get_resize_target(x, crop_target, do_crop=(resize_method==ResizeMethod.CROP))
x.resize(target)
elif resize_method==ResizeMethod.SQUISH: x.resize((x.shape[0],) + crop_target)
else: size = x.size
size_tfms = [o for o in tfms if isinstance(o.tfm,TfmCrop)]
for tfm in tfms:
if tfm.tfm in xtra: x = tfm(x, **xtra[tfm.tfm])
elif tfm in size_tfms:
if resize_method in (ResizeMethod.CROP,ResizeMethod.PAD):
x = tfm(x, size=_get_crop_target(size,mult=mult), padding_mode=padding_mode)
else: x = tfm(x)
return x.refresh()
def refresh(self)->None:
"Apply any logit, flow, or affine transfers that have been sent to the `Image`."
if self._logit_px is not None:
self._px = self._logit_px.sigmoid_()
self._logit_px = None
if self._affine_mat is not None or self._flow is not None:
self._px = _grid_sample(self._px, self.flow, **self.sample_kwargs)
self.sample_kwargs = {}
self._flow = None
return self
def save(self, fn:PathOrStr):
"Save the image to `fn`."
x = image2np(self.data*255).astype(np.uint8)
PIL.Image.fromarray(x).save(fn)
@property
def px(self)->TensorImage:
"Get the tensor pixel buffer."
self.refresh()
return self._px
@px.setter
def px(self,v:TensorImage)->None:
"Set the pixel buffer to `v`."
self._px=v
@property
def flow(self)->FlowField:
"Access the flow-field grid after applying queued affine transforms."
if self._flow is None:
self._flow = _affine_grid(self.shape)
if self._affine_mat is not None:
self._flow = _affine_mult(self._flow,self._affine_mat)
self._affine_mat = None
return self._flow
@flow.setter
def flow(self,v:FlowField): self._flow=v
def lighting(self, func:LightingFunc, *args:Any, **kwargs:Any):
"Equivalent to `image = sigmoid(func(logit(image)))`."
self.logit_px = func(self.logit_px, *args, **kwargs)
return self
def pixel(self, func:PixelFunc, *args, **kwargs)->'Image':
"Equivalent to `image.px = func(image.px)`."
self.px = func(self.px, *args, **kwargs)
return self
def coord(self, func:CoordFunc, *args, **kwargs)->'Image':
"Equivalent to `image.flow = func(image.flow, image.size)`."
self.flow = func(self.flow, *args, **kwargs)
return self
def affine(self, func:AffineFunc, *args, **kwargs)->'Image':
"Equivalent to `image.affine_mat = image.affine_mat @ func()`."
m = tensor(func(*args, **kwargs)).to(self.device)
self.affine_mat = self.affine_mat @ m
return self
def resize(self, size:Union[int,TensorImageSize])->'Image':
"Resize the image to `size`, size can be a single int."
assert self._flow is None
if isinstance(size, int): size=(self.shape[0], size, size)
if tuple(size)==tuple(self.shape): return self
self.flow = _affine_grid(size)
return self
@property
def affine_mat(self)->AffineMatrix:
"Get the affine matrix that will be applied by `refresh`."
if self._affine_mat is None:
self._affine_mat = torch.eye(3).to(self.device)
return self._affine_mat
@affine_mat.setter
def affine_mat(self,v)->None: self._affine_mat=v
@property
def logit_px(self)->LogitTensorImage:
"Get logit(image.px)."
if self._logit_px is None: self._logit_px = logit_(self.px)
return self._logit_px
@logit_px.setter
def logit_px(self,v:LogitTensorImage)->None: self._logit_px=v
@property
def data(self)->TensorImage:
"Return this images pixels as a tensor."
return self.px
def show(self, ax:plt.Axes=None, figsize:tuple=(3,3), title:Optional[str]=None, hide_axis:bool=True,
cmap:str=None, y:Any=None, **kwargs):
"Show image on `ax` with `title`, using `cmap` if single-channel, overlaid with optional `y`"
cmap = ifnone(cmap, defaults.cmap)
ax = show_image(self, ax=ax, hide_axis=hide_axis, cmap=cmap, figsize=figsize)
if y is not None: y.show(ax=ax, **kwargs)
if title is not None: ax.set_title(title)
The code can be found in fastai.vision.image