I am using carnava dataset that has images of size (1280, 1918, 3) deeplabv3, with resnet101 as backbone. The batch size is 2, since the imges are really large.
Here is the complete code.
class carnava(Dataset):
def __init__(self,path,transform):
self.path = path
self.PATH2 =os.path.join(self.path,'train_masks_png')
img_list = [img for img in os.listdir(os.path.join(self.path,'train'))]
mask_list = [mask for mask in os.listdir(os.path.join(self.path,'train_masks'))]
mask_png_list = [Image.open(os.path.join(self.path,'train_masks',ms)).save(os.path.join(self.PATH2,f'{ms[:-4]}.png'))
for ms in mask_list]
self.img_list= img_list
self.mask_list= mask_list
self.mask_png_list = mask_png_list
self.transform = transform
self.len= len(img_list)
def __getitem__(self,index):
img = Image.open(os.path.join(self.path,'train',self.img_list[index]))
mask = Image.open(os.path.join(self.PATH2,f'{self.mask_list[index][:-4]}.png'))
data = {'image':np.array(img),'mask':np.array(mask)}
data_tr = self.transform(**data)
img_tr = data_tr['image']
mask_tr = data_tr['mask']
return(img_tr,mask_tr)
def __len__(self):
return(self.len)
train_loader= DataLoader(train_set, batch_size =2, shuffle=True,num_workers=4)
val_loader = DataLoader(val_set,batch_size=4, shuffle=False, num_workers=4)
model = deeplabv3_resnet101(pretrained=True, progress=True)
model.classifier[4]= nn.Conv2d(in_channels=256,out_channels=1,kernel_size=1,stride=1)
def dice_loss(pred, target):
smooth = 1.
iflat = pred.contiguous().view(-1)
tflat = target.contiguous().view(-1)
intersection = (iflat * tflat).sum()
A_sum = torch.sum(iflat * iflat)
B_sum = torch.sum(tflat * tflat)
return 1 - ((2. * intersection + smooth) / (A_sum + B_sum + smooth) )
def train(epoch,model,optimizer,criterion,device,phase):
if phase=='train':
model.train()
dataloader = train_loader
else:
model.eval()
dataloader = val_loader
for i, (img, mask) in enumerate(dataloader):
img = img.to(device)
mask = mask.to(device)
target = model(img)['out']
loss = criterion(target,mask)
loss_total += loss.item()
if phase=='train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
target=target.detach().to(device)
dis_loss += dice_loss(target,mask)
return(loss_total/len(train_loader), dis_loss/len(train_loader))
def main(num_epochs):
device= torch.device(5)
model = get_model('deeplab').to(device)
optimizer = optim.Adam(model.parameters(),lr=0.1)
criterion =torch.nn.BCEWithLogitsLoss().to(device)
for epoch in range(num_epochs):
train_loss,dic_loss_t = train(epoch,model,optimizer,criterion,device,'train')
val_loss, dis_loss_v = train(epoch,odel,optimizer,criterion,device,'val')
if __name__=="__main__":
main(10)
CUDA out of memory. Tried to allocate 2.00 MiB (GPU 5; 10.76 GiB total capacity; 112.99 MiB already allocated; 5.25 MiB free; 118.00 MiB reserved in total by PyTorch)