def __init__(self,
image_size: int,
mask_channels_count: int,
image_channels_count: int = 3,
generator_size: int = 32,
discriminator_size: int = 32):
super().__init__()
# netG = UNetGenerator(noise, image_size, mask_channels_count, image_channels_count, generator_size) \
netG = ResnetGenerator(mask_channels_count, image_channels_count, generator_size, n_blocks=9)\
.to(ParallelConfig.MAIN_DEVICE)
# netD = Discriminator(discriminator_size, image_channels_count + mask_channels_count, image_size) \
netD = ResDiscriminator(image_channels_count + mask_channels_count, discriminator_size, img_f=256, layers=4) \
.to(ParallelConfig.MAIN_DEVICE)
if torch.cuda.device_count() > 1:
netD = nn.DataParallel(netD, ParallelConfig.GPU_IDS)
netG = nn.DataParallel(netG, ParallelConfig.GPU_IDS)
self.gan_model = ConditionalGANModel(
netG,
HingeLoss(netD) + GANLossObject(
lambda x, y: Loss.ZERO(),
lambda dgz, real, fake: Loss(nn.L1Loss()(fake[0], real[0])) * 0.1,
netD
)
)
def train(self, image: Tensor, sp: Tensor) -> Loss:
segm = self.segmentation(image)
loss: Loss = Loss.ZERO()
for pen in self.penalties:
loss = loss + pen.forward(image, sp, segm)
loss.minimize_step(self.opt)
return loss
def loss_backward(self, condition: Dict[str, Tensor]) -> Loss:
condition_pred: Dict[str, Tensor] = self.g_forward(
self.g_backward(condition))
loss = Loss.ZERO()
for name in condition.keys():
loss += self.loss_2[name](condition_pred[name], condition[name])
return loss
def loss_forward(self, condition: Dict[str, T1]) -> Loss:
t2 = self.g_forward(condition)
condition_pred: Dict[str, T1] = self.g_backward(t2)
loss = Loss.ZERO()
for name in condition.keys():
loss += self.loss_1[name](condition_pred[name], condition[name])
return loss
def discriminator_loss(self, d_real: Tensor, d_fake: Tensor) -> Loss:
return Loss.ZERO()