def _make_images_board(self, model):
model.eval()
num_imgs = 64
batch = next(iter(self.data_loaders[1]))
input_images, depthGT, maskGT = utils.unpack_batch_fixed(
batch, self.cfg.device)
with torch.set_grad_enabled(False):
XYZ, maskLogit = model(input_images)
XY = XYZ[:, :self.cfg.outViewN * 2, :, :]
depth = XYZ[:, self.cfg.outViewN * 2:self.cfg.outViewN * 3, :, :]
mask = (maskLogit > 0).float()
return {
'RGB':
utils.make_grid(input_images[:num_imgs]),
'depth':
utils.make_grid(1 - depth[:num_imgs, 0:1, :, :]),
'depth_mask':
utils.make_grid(((1 - depth) * mask)[:num_imgs, 0:1, :, :]),
'depthGT':
utils.make_grid(1 - depthGT[:num_imgs, 0:1, :, :]),
'mask':
utils.make_grid(torch.sigmoid(maskLogit[:num_imgs, 0:1, :, :])),
'maskGT':
utils.make_grid(maskGT[:num_imgs, 0:1, :, :]),
}
def findLR(self, model, optimizer, writer,
start_lr=1e-7, end_lr=10, num_iters=50):
model.train()
losses = []
lrs = np.logspace(np.log10(start_lr), np.log10(end_lr), num_iters)
for lr in lrs:
# Update LR
for group in optimizer.param_groups: group['lr'] = lr
batch = next(iter(self.data_loaders[0]))
input_images, depthGT, maskGT = utils.unpack_batch_fixed(batch, self.cfg.device)
# ------ define ground truth------
XGT, YGT = torch.meshgrid([torch.arange(self.cfg.outH), # [H,W]
torch.arange(self.cfg.outW)]) # [H,W]
XGT, YGT = XGT.float(), YGT.float()
XYGT = torch.cat([
XGT.repeat([self.cfg.outViewN, 1, 1]),
YGT.repeat([self.cfg.outViewN, 1, 1])], dim=0) #[2V,H,W]
XYGT = XYGT.unsqueeze(dim=0).to(self.cfg.device) #[1,2V,H,W]
with torch.set_grad_enabled(True):
optimizer.zero_grad()
XYZ, maskLogit = model(input_images)
XY = XYZ[:, :self.cfg.outViewN * 2, :, :]
depth = XYZ[:, self.cfg.outViewN * 2:self.cfg.outViewN * 3, :, :]
mask = (maskLogit > 0).byte()
# ------ Compute loss ------
loss_XYZ = self.l1(XY, XYGT)
loss_XYZ += self.l1(depth.masked_select(mask),
depthGT.masked_select(mask))
loss_mask = self.sigmoid_bce(maskLogit, maskGT)
loss = loss_mask + self.cfg.lambdaDepth * loss_XYZ
# Update weights
loss.backward()
# True Weight decay
if self.cfg.trueWD is not None:
for group in optimizer.param_groups:
for param in group['params']:
param.data = param.data.add(
-self.cfg.trueWD * group['lr'], param.data)
optimizer.step()
losses.append(loss.item())
fig, ax = plt.subplots()
ax.plot(lrs, losses)
ax.set_xlabel('learning rate')
ax.set_ylabel('loss')
ax.set_xscale('log')
writer.add_figure('findLR', fig)
def findLR(self, model, optimizer, writer,
start_lr=1e-7, end_lr=10, num_iters=50):
model.train()
losses = []
lrs = np.logspace(np.log10(start_lr), np.log10(end_lr), num_iters)
for lr in lrs:
# Update LR
for group in optimizer.param_groups: group['lr'] = lr
batch = next(iter(self.data_loaders[0]))
input_images, depthGT, maskGT = utils.unpack_batch_fixed(batch, self.cfg.device)
# ------ define ground truth------
XGT, YGT = torch.meshgrid([torch.arange(self.cfg.outH), # [H,W]
torch.arange(self.cfg.outW)]) # [H,W]
XGT, YGT = XGT.float(), YGT.float()
XYGT = torch.cat([
XGT.repeat([self.cfg.outViewN, 1, 1]),
YGT.repeat([self.cfg.outViewN, 1, 1])], dim=0) #[2V,H,W]
XYGT = XYGT.unsqueeze(dim=0).to(self.cfg.device) #[1,2V,H,W]
with torch.set_grad_enabled(True):
optimizer.zero_grad()
XYZ, maskLogit = model(input_images)
XY = XYZ[:, :self.cfg.outViewN * 2, :, :]
depth = XYZ[:, self.cfg.outViewN * 2:self.cfg.outViewN * 3, :, :]
mask = (maskLogit > 0).byte()
# ------ Compute loss ------
loss_XYZ = self.l1(XY, XYGT)
loss_XYZ += self.l1(depth.masked_select(mask),
depthGT.masked_select(mask))
loss_mask = self.sigmoid_bce(maskLogit, maskGT)
loss = loss_mask + self.cfg.lambdaDepth * loss_XYZ
# Update weights
loss.backward()
# True Weight decay
if self.cfg.trueWD is not None:
for group in optimizer.param_groups:
for param in group['params']:
param.data = param.data.add(
-self.cfg.trueWD * group['lr'], param.data)
optimizer.step()
losses.append(loss.item())
fig, ax = plt.subplots()
ax.plot(lrs, losses)
ax.set_xlabel('learning rate')
ax.set_ylabel('loss')
ax.set_xscale('log')
writer.add_figure('findLR', fig)
def _val_on_epoch(self, model):
model.eval()
data_loader = self.data_loaders[1]
running_loss_XYZ = 0.0
running_loss_mask = 0.0
running_loss = 0.0
for batch in data_loader:
input_images, depthGT, maskGT = utils.unpack_batch_fixed(
batch, self.cfg.device)
# ------ define ground truth------
XGT, YGT = torch.meshgrid([
torch.arange(self.cfg.outH), # [H,W]
torch.arange(self.cfg.outW)
]) # [H,W]
XGT, YGT = XGT.float(), YGT.float()
XYGT = torch.cat([
XGT.repeat([self.cfg.outViewN, 1, 1]),
YGT.repeat([self.cfg.outViewN, 1, 1])
],
dim=0) #[2V,H,W]
XYGT = XYGT.unsqueeze(dim=0).to(self.cfg.device) # [1,2V,H,W]
with torch.set_grad_enabled(False):
XYZ, maskLogit = model(input_images)
XY = XYZ[:, :self.cfg.outViewN * 2, :, :]
depth = XYZ[:,
self.cfg.outViewN * 2:self.cfg.outViewN * 3, :, :]
mask = (maskLogit > 0).byte()
# ------ Compute loss ------
loss_XYZ = self.l1(XY, XYGT)
loss_XYZ += self.l1(depth.masked_select(mask),
depthGT.masked_select(mask))
loss_mask = self.sigmoid_bce(maskLogit, maskGT)
loss = loss_mask + self.cfg.lambdaDepth * loss_XYZ
running_loss_XYZ += loss_XYZ.item() * input_images.size(0)
running_loss_mask += loss_mask.item() * input_images.size(0)
running_loss += loss.item() * input_images.size(0)
epoch_loss_XYZ = running_loss_XYZ / len(data_loader.dataset)
epoch_loss_mask = running_loss_mask / len(data_loader.dataset)
epoch_loss = running_loss / len(data_loader.dataset)
print(f"\tVal loss: {epoch_loss}")
return {
"epoch_loss_XYZ": epoch_loss_XYZ,
"epoch_loss_mask": epoch_loss_mask,
"epoch_loss": epoch_loss,
}
def _val_on_epoch(self, model):
model.eval()
data_loader = self.data_loaders[1]
running_loss_XYZ = 0.0
running_loss_mask = 0.0
running_loss = 0.0
for batch in data_loader:
input_images, depthGT, maskGT = utils.unpack_batch_fixed(batch, self.cfg.device)
# ------ define ground truth------
XGT, YGT = torch.meshgrid([
torch.arange(self.cfg.outH), # [H,W]
torch.arange(self.cfg.outW)]) # [H,W]
XGT, YGT = XGT.float(), YGT.float()
XYGT = torch.cat([
XGT.repeat([self.cfg.outViewN, 1, 1]),
YGT.repeat([self.cfg.outViewN, 1, 1])], dim=0) #[2V,H,W]
XYGT = XYGT.unsqueeze(dim=0).to(self.cfg.device) # [1,2V,H,W]
with torch.set_grad_enabled(False):
XYZ, maskLogit = model(input_images)
XY = XYZ[:, :self.cfg.outViewN * 2, :, :]
depth = XYZ[:, self.cfg.outViewN * 2:self.cfg.outViewN*3,:,:]
mask = (maskLogit > 0).byte()
# ------ Compute loss ------
loss_XYZ = self.l1(XY, XYGT)
loss_XYZ += self.l1(depth.masked_select(mask),
depthGT.masked_select(mask))
loss_mask = self.sigmoid_bce(maskLogit, maskGT)
loss = loss_mask + self.cfg.lambdaDepth * loss_XYZ
running_loss_XYZ += loss_XYZ.item() * input_images.size(0)
running_loss_mask += loss_mask.item() * input_images.size(0)
running_loss += loss.item() * input_images.size(0)
epoch_loss_XYZ = running_loss_XYZ / len(data_loader.dataset)
epoch_loss_mask = running_loss_mask / len(data_loader.dataset)
epoch_loss = running_loss / len(data_loader.dataset)
print(f"\tVal loss: {epoch_loss}")
return {"epoch_loss_XYZ": epoch_loss_XYZ,
"epoch_loss_mask": epoch_loss_mask,
"epoch_loss": epoch_loss, }
def _make_images_board(self, model):
model.eval()
num_imgs = 64
batch = next(iter(self.data_loaders[1]))
input_images, depthGT, maskGT = utils.unpack_batch_fixed(batch, self.cfg.device)
with torch.set_grad_enabled(False):
XYZ, maskLogit = model(input_images)
XY = XYZ[:, :self.cfg.outViewN * 2, :, :]
depth = XYZ[:, self.cfg.outViewN * 2:self.cfg.outViewN * 3, :, :]
mask = (maskLogit > 0).float()
return {'RGB': utils.make_grid(input_images[:num_imgs]),
'depth': utils.make_grid(1-depth[:num_imgs, 0:1, :, :]),
'depth_mask': utils.make_grid(
((1-depth)*mask)[:num_imgs, 0:1, :, :]),
'depthGT': utils.make_grid(
1-depthGT[:num_imgs, 0:1, :, :]),
'mask': utils.make_grid(
torch.sigmoid(maskLogit[:num_imgs, 0:1,:, :])),
'maskGT': utils.make_grid(maskGT[:num_imgs, 0:1, :, :]),
}
def _train_on_epoch(self, model, optimizer):
#model=nn.DataParallel(model,device_ids=[0,1])
model.train()
data_loader = self.data_loaders[0]
running_loss_XYZ = 0.0
running_loss_mask = 0.0
running_loss = 0.0
for self.iteration, batch in enumerate(data_loader, self.iteration):
input_images, depthGT, maskGT = utils.unpack_batch_fixed(batch, self.cfg.device)
# ------ define ground truth------
XGT, YGT = torch.meshgrid([
torch.arange(self.cfg.outH), # [H,W]
torch.arange(self.cfg.outW)]) # [H,W]
XGT, YGT = XGT.float(), YGT.float()
XYGT = torch.cat([
XGT.repeat([self.cfg.outViewN, 1, 1]),
YGT.repeat([self.cfg.outViewN, 1, 1])], dim=0) #[2V,H,W]
XYGT = XYGT.unsqueeze(dim=0).to(self.cfg.device) # [1,2V,H,W]
with torch.set_grad_enabled(True):
optimizer.zero_grad()
XYZ, maskLogit = model(input_images)
XY = XYZ[:, :self.cfg.outViewN * 2, :, :]
depth = XYZ[:, self.cfg.outViewN * 2:self.cfg.outViewN * 3, :, :]
####################################################################
##################################
tmp_1 = torch.cat([maskLogit[:,0:1,:,:],maskLogit[:,2:3,:,:],maskLogit[:,4:5,:,:],
maskLogit[:,6:7,:,:],maskLogit[:,7:8,:,:],maskLogit[:,9:10,:,:],
maskLogit[:,11:12,:,:],maskLogit[:,13:14,:,:]],1)
#print(tmp_1.size())
tmp_2 = torch.cat([maskLogit[:,1:2,:,:],maskLogit[:,3:4,:,:],maskLogit[:,5:6,:,:],
maskLogit[:,7:8,:,:],maskLogit[:,9:10,:,:],maskLogit[:,11:12,:,:],
maskLogit[:,13:14,:,:],maskLogit[:,15:16,:,:]],1)
#mask = (maskLogit > 0).byte()
mask = (tmp_2 > 0).byte()
###################################
# print(mask.type())
# print(mask.size())
# print(maskGT.type())
# print(maskGT.size())
# ------ Compute loss ------
loss_XYZ = self.l1(XY, XYGT)
####################################################################
loss_XYZ += self.l1(depth.masked_select(mask),
depthGT.masked_select(mask))
####################################################################
#######################################################
#loss_mask = self.sigmoid_bce(maskLogit, maskGT.long())
# print(maskLogit[:,0:8,:,:].type())
# print(maskLogit[:,0:8,:,:].size())
tmp_1 = torch.cat([maskLogit[:,0:1,:,:],maskLogit[:,2:3,:,:],maskLogit[:,4:5,:,:],
maskLogit[:,6:7,:,:],maskLogit[:,7:8,:,:],maskLogit[:,9:10,:,:],
maskLogit[:,11:12,:,:],maskLogit[:,13:14,:,:]],1)
#print(tmp_1.size())
tmp_2 = torch.cat([maskLogit[:,1:2,:,:],maskLogit[:,3:4,:,:],maskLogit[:,5:6,:,:],
maskLogit[:,7:8,:,:],maskLogit[:,9:10,:,:],maskLogit[:,11:12,:,:],
maskLogit[:,13:14,:,:],maskLogit[:,15:16,:,:]],1)
#print(tmp_2.size())
#maskLogit = torch.cat([tmp_1,tmp_2],4)
maskLogit = torch.stack([tmp_1,tmp_2],dim=1)
#print(maskLogit.size())
#maskLogit = torch.stack([maskLogit[:,0:8,:,:],maskLogit[:,8:16,:,:]],dim=1)
# print(maskLogit.type())
# print(maskLogit.size())
loss_mask = self.cross_entropy(maskLogit, maskGT.long())
#######################################################
#loss_mask = self.sigmoid_bce(maskLogit, maskGT)
loss = loss_mask + self.cfg.lambdaDepth * loss_XYZ
# ------ Update weights ------
loss.backward()
# True Weight decay
if self.cfg.trueWD is not None:
for group in optimizer.param_groups:
for param in group['params']:
param.data.add_(
-self.cfg.trueWD * group['lr'], param.data)
optimizer.step()
if self.on_after_batch is not None:
if self.cfg.lrSched.lower() in "cyclical":
self.on_after_batch(self.iteration)
else: self.on_after_batch(self.epoch)
running_loss_XYZ += loss_XYZ.item() * input_images.size(0)
running_loss_mask += loss_mask.item() * input_images.size(0)
running_loss += loss.item() * input_images.size(0)
epoch_loss_XYZ = running_loss_XYZ / len(data_loader.dataset)
epoch_loss_mask = running_loss_mask / len(data_loader.dataset)
epoch_loss = running_loss / len(data_loader.dataset)
print(f"\tTrain loss: {epoch_loss}")
return {"epoch_loss_XYZ": epoch_loss_XYZ,
"epoch_loss_mask": epoch_loss_mask,
"epoch_loss": epoch_loss, }
def _val_on_epoch(self, model):
model.eval()
data_loader = self.data_loaders[1]
running_loss_XYZ = 0.0
running_loss_mask = 0.0
running_loss = 0.0
for batch in data_loader:
input_images, depthGT, maskGT = utils.unpack_batch_fixed(batch, self.cfg.device)
# ------ define ground truth------
XGT, YGT = torch.meshgrid([
torch.arange(self.cfg.outH), # [H,W]
torch.arange(self.cfg.outW)]) # [H,W]
XGT, YGT = XGT.float(), YGT.float()
XYGT = torch.cat([
XGT.repeat([self.cfg.outViewN, 1, 1]),
YGT.repeat([self.cfg.outViewN, 1, 1])], dim=0) #[2V,H,W]
XYGT = XYGT.unsqueeze(dim=0).to(self.cfg.device) # [1,2V,H,W]
with torch.set_grad_enabled(False):
XYZ, maskLogit = model(input_images)
XY = XYZ[:, :self.cfg.outViewN * 2, :, :]
depth = XYZ[:, self.cfg.outViewN * 2:self.cfg.outViewN*3,:,:]
#######################################################
tmp_1 = torch.cat([maskLogit[:,0:1,:,:],maskLogit[:,2:3,:,:],maskLogit[:,4:5,:,:],
maskLogit[:,6:7,:,:],maskLogit[:,7:8,:,:],maskLogit[:,9:10,:,:],
maskLogit[:,11:12,:,:],maskLogit[:,13:14,:,:]],1)
#print(tmp_1.size())
tmp_2 = torch.cat([maskLogit[:,1:2,:,:],maskLogit[:,3:4,:,:],maskLogit[:,5:6,:,:],
maskLogit[:,7:8,:,:],maskLogit[:,9:10,:,:],maskLogit[:,11:12,:,:],
maskLogit[:,13:14,:,:],maskLogit[:,15:16,:,:]],1)
##################################
#mask = (maskLogit > 0).byte()
#mask = (maskLogit > 0).byte()
mask = (tmp_2 > 0).byte()
###################################
# ------ Compute loss ------
loss_XYZ = self.l1(XY, XYGT)
loss_XYZ += self.l1(depth.masked_select(mask),
depthGT.masked_select(mask))
#######################################################
#loss_mask = self.sigmoid_bce(maskLogit, maskGT.long())
# print(maskLogit[:,0:8,:,:].type())
# print(maskLogit[:,0:8,:,:].size())
tmp_1 = torch.cat([maskLogit[:,0:1,:,:],maskLogit[:,2:3,:,:],maskLogit[:,4:5,:,:],
maskLogit[:,6:7,:,:],maskLogit[:,7:8,:,:],maskLogit[:,9:10,:,:],
maskLogit[:,11:12,:,:],maskLogit[:,13:14,:,:]],1)
#print(tmp_1.size())
tmp_2 = torch.cat([maskLogit[:,1:2,:,:],maskLogit[:,3:4,:,:],maskLogit[:,5:6,:,:],
maskLogit[:,7:8,:,:],maskLogit[:,9:10,:,:],maskLogit[:,11:12,:,:],
maskLogit[:,13:14,:,:],maskLogit[:,15:16,:,:]],1)
#print(tmp_2.size())
#maskLogit = torch.cat([tmp_1,tmp_2],4)
maskLogit = torch.stack([tmp_1,tmp_2],dim=1)
#maskLogit = torch.stack([maskLogit[:,0:8,:,:],maskLogit[:,8:16,:,:]],dim=1)
# print(maskLogit.type())
# print(maskLogit.size())
loss_mask = self.cross_entropy(maskLogit, maskGT.long())
#######################################################
#loss_mask = self.sigmoid_bce(maskLogit, maskGT)
#loss_mask = self.sigmoid_bce(maskLogit, maskGT)
loss = loss_mask + self.cfg.lambdaDepth * loss_XYZ
running_loss_XYZ += loss_XYZ.item() * input_images.size(0)
running_loss_mask += loss_mask.item() * input_images.size(0)
running_loss += loss.item() * input_images.size(0)
epoch_loss_XYZ = running_loss_XYZ / len(data_loader.dataset)
epoch_loss_mask = running_loss_mask / len(data_loader.dataset)
epoch_loss = running_loss / len(data_loader.dataset)
print(f"\tVal loss: {epoch_loss}")
return {"epoch_loss_XYZ": epoch_loss_XYZ,
"epoch_loss_mask": epoch_loss_mask,
"epoch_loss": epoch_loss, }
def _train_on_epoch(self, model, optimizer):
model.train()
data_loader = self.data_loaders[0]
running_loss_XYZ = 0.0
running_loss_mask = 0.0
running_loss = 0.0
for self.iteration, batch in enumerate(data_loader, self.iteration):
input_images, depthGT, maskGT = utils.unpack_batch_fixed(
batch, self.cfg.device)
# ------ define ground truth------
XGT, YGT = torch.meshgrid([
torch.arange(self.cfg.outH), # [H,W]
torch.arange(self.cfg.outW)
]) # [H,W]
XGT, YGT = XGT.float(), YGT.float()
XYGT = torch.cat([
XGT.repeat([self.cfg.outViewN, 1, 1]),
YGT.repeat([self.cfg.outViewN, 1, 1])
],
dim=0) #[2V,H,W]
XYGT = XYGT.unsqueeze(dim=0).to(self.cfg.device) # [1,2V,H,W]
with torch.set_grad_enabled(True):
optimizer.zero_grad()
XYZ, maskLogit = model(input_images)
XY = XYZ[:, :self.cfg.outViewN * 2, :, :]
depth = XYZ[:,
self.cfg.outViewN * 2:self.cfg.outViewN * 3, :, :]
mask = (maskLogit > 0).byte()
# ------ Compute loss ------
loss_XYZ = self.l1(XY, XYGT)
loss_XYZ += self.l1(depth.masked_select(mask),
depthGT.masked_select(mask))
loss_mask = self.sigmoid_bce(maskLogit, maskGT)
loss = loss_mask + self.cfg.lambdaDepth * loss_XYZ
# ------ Update weights ------
loss.backward()
# True Weight decay
if self.cfg.trueWD is not None:
for group in optimizer.param_groups:
for param in group['params']:
param.data.add_(-self.cfg.trueWD * group['lr'],
param.data)
optimizer.step()
if self.on_after_batch is not None:
if self.cfg.lrSched.lower() in "cyclical":
self.on_after_batch(self.iteration)
else:
self.on_after_batch(self.epoch)
running_loss_XYZ += loss_XYZ.item() * input_images.size(0)
running_loss_mask += loss_mask.item() * input_images.size(0)
running_loss += loss.item() * input_images.size(0)
epoch_loss_XYZ = running_loss_XYZ / len(data_loader.dataset)
epoch_loss_mask = running_loss_mask / len(data_loader.dataset)
epoch_loss = running_loss / len(data_loader.dataset)
print(f"\tTrain loss: {epoch_loss}")
return {
"epoch_loss_XYZ": epoch_loss_XYZ,
"epoch_loss_mask": epoch_loss_mask,
"epoch_loss": epoch_loss,
}
def _train_on_epoch(self, model, optimizer):
model.train()
data_loader = self.data_loaders[0]
running_loss_XYZ = 0.0
running_loss_mask = 0.0
running_loss = 0.0
for self.iteration, batch in enumerate(data_loader, self.iteration):
input_images, depthGT, maskGT = utils.unpack_batch_fixed(batch, self.cfg.device)
# ------ define ground truth------
XGT, YGT = torch.meshgrid([
torch.arange(self.cfg.outH), # [H,W]
torch.arange(self.cfg.outW)]) # [H,W]
XGT, YGT = XGT.float(), YGT.float()
XYGT = torch.cat([
XGT.repeat([self.cfg.outViewN, 1, 1]),
YGT.repeat([self.cfg.outViewN, 1, 1])], dim=0) #[2V,H,W]
XYGT = XYGT.unsqueeze(dim=0).to(self.cfg.device) # [1,2V,H,W]
with torch.set_grad_enabled(True):
optimizer.zero_grad()
XYZ, maskLogit = model(input_images)
XY = XYZ[:, :self.cfg.outViewN * 2, :, :]
depth = XYZ[:, self.cfg.outViewN * 2:self.cfg.outViewN * 3, :, :]
mask = (maskLogit > 0).byte()
# ------ Compute loss ------
loss_XYZ = self.l1(XY, XYGT)
loss_XYZ += self.l1(depth.masked_select(mask),
depthGT.masked_select(mask))
loss_mask = self.sigmoid_bce(maskLogit, maskGT)
loss = loss_mask + self.cfg.lambdaDepth * loss_XYZ
# ------ Update weights ------
loss.backward()
# True Weight decay
if self.cfg.trueWD is not None:
for group in optimizer.param_groups:
for param in group['params']:
param.data.add_(
-self.cfg.trueWD * group['lr'], param.data)
optimizer.step()
if self.on_after_batch is not None:
if self.cfg.lrSched.lower() in "cyclical":
self.on_after_batch(self.iteration)
else: self.on_after_batch(self.epoch)
running_loss_XYZ += loss_XYZ.item() * input_images.size(0)
running_loss_mask += loss_mask.item() * input_images.size(0)
running_loss += loss.item() * input_images.size(0)
epoch_loss_XYZ = running_loss_XYZ / len(data_loader.dataset)
epoch_loss_mask = running_loss_mask / len(data_loader.dataset)
epoch_loss = running_loss / len(data_loader.dataset)
print(f"\tTrain loss: {epoch_loss}")
return {"epoch_loss_XYZ": epoch_loss_XYZ,
"epoch_loss_mask": epoch_loss_mask,
"epoch_loss": epoch_loss, }