def train(engine: Engine, ctx: Context):
datas, targets = ctx.inputs
ctx.datas, ctx.targets = datas.to(ctx.device), targets.to(ctx.device)
ctx.outputs = ctx.net(ctx.datas)
ctx.loss = ctx.criterion(ctx.outputs, ctx.targets)
ctx.optimizer.zero_grad()
ctx.loss.backward()
ctx.optimizer.step()
def train_discriminator(engine: Engine, ctx: GANContext):
ctx.d_optimizer.zero_grad()
ctx.real_datas, noisy_datas = ctx.inputs
d_real_outputs = ctx.discriminator(ctx.real_datas)
ctx.d_real_loss = ctx.criterion(d_real_outputs, ctx.ones)
flame.logger.debug(ctx.d_real_loss)
ctx.d_real_loss.backward()
ctx.d_fake_datas = ctx.generator(noisy_datas).detach()
d_fake_pridiction = ctx.discriminator(ctx.d_fake_datas)
ctx.d_fake_loss = ctx.criterion(d_fake_pridiction, ctx.zeros)
ctx.d_fake_loss.backward()
ctx.d_optimizer.step()
def train_generator(engine: Engine, ctx: GANContext):
ctx.g_optimizer.zero_grad()
noisy_datas = ctx.inputs
ctx.g_fake_data = ctx.generator(noisy_datas)
g_fake_pridiction = ctx.discriminator(ctx.g_fake_data)
ctx.g_loss = ctx.criterion(g_fake_pridiction, ctx.ones)
ctx.g_loss.backward()
ctx.g_optimizer.step()
def debug_train(engine: Engine, ctx: Context):
datas, targets = ctx.inputs
ctx.datas, ctx.targets = datas.to(ctx.device), targets.to(ctx.device)
ctx.timer.mark("to device")
ctx.outputs = ctx.net(ctx.datas)
ctx.timer.mark("forward")
ctx.loss = ctx.criterion(ctx.outputs, ctx.targets)
ctx.timer.mark("criterion")
ctx.optimizer.zero_grad()
ctx.loss.backward()
ctx.timer.mark("backward")
ctx.optimizer.step()
ctx.timer.mark("step")
def validation(engine: Engine, ctx: Context):
datas, targets = ctx.inputs
ctx.datas, ctx.targets = datas.to(ctx.device), targets.to(ctx.device)
ctx.outputs = ctx.net(ctx.datas)
ctx.loss = ctx.criterion(ctx.outputs, ctx.targets)