ZHKKKe / MODNet

A Trimap-Free Portrait Matting Solution in Real Time [AAAI 2022]
Apache License 2.0
3.83k stars 637 forks source link

runtime error in training on single image #75

Closed Randheerkumar closed 3 years ago

Randheerkumar commented 3 years ago

getting this: RuntimeError: Sizes of tensors must match except in dimension 2. Got 87 and 88 (The offending index is 0)

code is: import numpy as np from src.models.modnet import MODNet from src.trainer import supervised_training_iter import torchvision.transforms.functional as TF

import cv2

bs = 1 # batch size lr = 0.01 # learn rate epochs = 40 # total epochs

modnet = torch.nn.DataParallel(MODNet()).cuda() optimizer = torch.optim.SGD(modnet.parameters(), lr=lr, momentum=0.9) lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(0.25 * epochs), gamma=0.1)

image=Image.open("image/2.jpg") #input image trimap=Image.open("trimap/trimap.png") # trimap gt_matte=Image.open("gt/2.png") #ground truth

image = TF.totensor(image) image.unsqueeze(0) print(image.shape) trimap = TF.totensor(trimap) trimap.unsqueeze(0) gt_matte = TF.to_tensor(gt_matte) gtmatte.unsqueeze(0)

dataloader = CREATE_YOUR_DATALOADER(bs) # NOTE: please finish this function

for epoch in range(0, epochs):

for idx, (image, trimap, gt_matte) in enumerate(dataloader):

semantic_loss, detail_loss, matte_loss = supervised_training_iter(modnet, optimizer, image, trimap, gt_matte)
#lr_scheduler.step()

@ZHKKKe

ZHKKKe commented 3 years ago

Hi, thanks for your attention. Please try to resize the input image to make sure its width/height can be divided by 32.