Skip to content
Snippets Groups Projects
Commit 730f3623 authored by Tamino Huxohl's avatar Tamino Huxohl
Browse files

replace old cgan with new cgan

parent 7fd7c7b2
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
mu_map/training/cgan.py
+ 137
− 183
View file @ 730f3623
from dataclasses import dataclass
import os
from typing import Dict, Optional
import sys
from logging import Logger
from typing import Optional
import torch
from torch import Tensor
from mu_map.dataset.default import MuMapDataset
from mu_map.training.lib import TrainingParams, AbstractTraining
from mu_map.training.loss import WeightedLoss
from mu_map.logging import get_logger
# Establish convention for real and fake labels during training
LABEL_REAL = 1.0
LABEL_FAKE = 0.0
@dataclass
class TrainingParams:
model: torch.nn.Module
optimizer: torch.optim.Optimizer
lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler]
class DiscriminatorParams(TrainingParams):
"""
Wrap training parameters to always carry the name 'Discriminator'.
"""
def __init__(
self,
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
):
super().__init__(
name="Discriminator",
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
)
class GeneratorParams(TrainingParams):
"""
Wrap training parameters to always carry the name 'Generator'.
"""
def __init__(
self,
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
):
super().__init__(
name="Generator",
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
)
class cGANTraining:
class cGANTraining(AbstractTraining):
"""
Implementation of a conditional generative adversarial network training.
"""
def __init__(
self,
data_loaders: Dict[str, torch.utils.data.DataLoader],
epochs: int,
dataset: MuMapDataset,
batch_size: int,
device: torch.device,
snapshot_dir: str,
snapshot_epoch: int,
params_generator: torch.nn.Module,
params_discriminator: torch.nn.Module,
params_generator: GeneratorParams,
params_discriminator: DiscriminatorParams,
loss_func_dist: WeightedLoss,
weight_criterion_dist: float,
weight_criterion_adv: float,
logger=None,
logger: Optional[Logger] = None,
):
self.data_loaders = data_loaders
self.epochs = epochs
self.device = device
"""
:param params_generator: training parameters containing a model an according optimizer and optionally a learning rate scheduler for the generator
:param params_discriminator: training parameters containing a model an according optimizer and optionally a learning rate scheduler for the discriminator
:param loss_func_dist: distance loss function for the generator
:param weight_criterion_dist: weight of the distance loss when training the generator
:param weight_criterion_adv: weight of the adversarial loss when training the generator
"""
super().__init__(
epochs, dataset, batch_size, device, snapshot_dir, snapshot_epoch, logger
)
self.training_params.append(params_generator)
self.training_params.append(params_discriminator)
self.snapshot_dir = snapshot_dir
self.snapshot_epoch = snapshot_epoch
self.logger = logger if logger is not None else get_logger(name=cGANTraining.__name__)
self.generator = params_generator.model
self.discriminator = params_discriminator.model
self.params_g = params_generator
self.params_d = params_discriminator
self.optim_g = params_generator.optimizer
self.optim_d = params_discriminator.optimizer
self.weight_criterion_dist = weight_criterion_dist
self.weight_criterion_adv = weight_criterion_adv
......@@ -53,134 +93,66 @@ class cGANTraining:
self.criterion_adv = torch.nn.MSELoss(reduction="mean")
self.criterion_dist = loss_func_dist
def run(self):
loss_val_min = sys.maxsize
for epoch in range(1, self.epochs + 1):
str_epoch = f"{str(epoch):>{len(str(self.epochs))}}"
self.logger.debug(f"Run epoch {str_epoch}/{self.epochs} ...")
self._train_epoch()
loss_train = self._eval_epoch("train")
self.logger.info(
f"Epoch {str_epoch}/{self.epochs} - Loss train: {loss_train:.6f}"
)
loss_val = self._eval_epoch("validation")
self.logger.info(
f"Epoch {str_epoch}/{self.epochs} - Loss validation: {loss_val:.6f}"
)
if loss_val < loss_val_min:
loss_val_min = loss_val
self.logger.info(
f"Store snapshot val_min of epoch {str_epoch} with minimal validation loss"
)
self.store_snapshot("val_min")
if epoch % self.snapshot_epoch == 0:
self._store_snapshot(epoch)
if self.params_d.lr_scheduler is not None:
self.params_d.lr_scheduler.step()
if self.params_g.lr_scheduler is not None:
self.params_g.lr_scheduler.step()
return loss_val_min
def _train_epoch(self):
# setup training mode
torch.set_grad_enabled(True)
self.params_d.model.train()
self.params_g.model.train()
data_loader = self.data_loaders["train"]
for i, (recons, mu_maps_real) in enumerate(data_loader):
print(
f"Batch {str(i):>{len(str(len(data_loader)))}}/{len(data_loader)}",
end="\r",
)
batch_size = recons.shape[0]
recons = recons.to(self.device)
mu_maps_real = mu_maps_real.to(self.device)
self.params_d.optimizer.zero_grad()
self.params_g.optimizer.zero_grad()
# compute fake mu maps with generator
mu_maps_fake = self.params_g.model(recons)
# compute discriminator loss for fake mu maps
inputs_d_fake = torch.cat((recons, mu_maps_fake), dim=1)
outputs_d_fake = self.params_d.model(
inputs_d_fake.detach()
) # note the detach, so that gradients are not computed for the generator
labels_fake = torch.full(
(outputs_d_fake.shape), LABEL_FAKE, device=self.device
)
loss_d_fake = self.criterion_adv(outputs_d_fake, labels_fake)
# compute discriminator loss for real mu maps
inputs_d_real = torch.cat((recons, mu_maps_real), dim=1)
outputs_d_real = self.params_d.model(
inputs_d_real
) # note the detach, so that gradients are not computed for the generator
labels_real = torch.full(
(outputs_d_fake.shape), LABEL_REAL, device=self.device
)
loss_d_real = self.criterion_adv(outputs_d_real, labels_real)
# update discriminator
loss_d = 0.5 * (loss_d_fake + loss_d_real)
loss_d.backward() # compute gradients
self.params_d.optimizer.step()
# update generator
inputs_d_fake = torch.cat((recons, mu_maps_fake), dim=1)
outputs_d_fake = self.params_d.model(inputs_d_fake)
loss_g_adv = self.criterion_adv(outputs_d_fake, labels_real)
loss_g_dist = self.criterion_dist(mu_maps_fake, mu_maps_real)
loss_g = (
self.weight_criterion_adv * loss_g_adv
+ self.weight_criterion_dist * loss_g_dist
)
loss_g.backward()
self.params_g.optimizer.step()
def _eval_epoch(self, split_name):
# setup evaluation mode
torch.set_grad_enabled(False)
self.params_d.model = self.params_d.model.eval()
self.params_g.model = self.params_g.model.eval()
data_loader = self.data_loaders[split_name]
loss = 0.0
updates = 0
for i, (recons, mu_maps) in enumerate(data_loader):
print(
f"Batch {str(i):>{len(str(len(data_loader)))}}/{len(data_loader)}",
end="\r",
)
recons = recons.to(self.device)
mu_maps = mu_maps.to(self.device)
outputs = self.params_g.model(recons)
loss += torch.nn.functional.l1_loss(outputs, mu_maps)
updates += 1
return loss / updates
def _after_train_batch(self):
"""
Overwrite calling step on all optimizers as this needs to be done
separately for the generator and discriminator during the training of
a batch.
"""
pass
def _train_batch(self, recons: torch.Tensor, mu_maps: torch.Tensor) -> float:
mu_maps_real = mu_maps # rename real mu maps for clarification
# compute fake mu maps with generator
mu_maps_fake = self.generator(recons)
# prepare inputs for the discriminator
inputs_d_fake = torch.cat((recons, mu_maps_fake), dim=1)
inputs_d_real = torch.cat((recons, mu_maps_real), dim=1)
# prepare labels/targets for the discriminator
labels_fake = torch.full(self.discriminator.get_output_shape(inputs_d_fake.shape), LABEL_FAKE, device=self.device)
labels_real = torch.full(self.discriminator.get_output_shape(inputs_d_real.shape), LABEL_REAL, device=self.device)
# ======================= Discriminator =====================================
# compute discriminator loss for fake mu maps
# detach is called so that gradients are not computed for the generator
outputs_d_fake = self.discriminator(inputs_d_fake.detach())
loss_d_fake = self.criterion_adv(outputs_d_fake, labels_fake)
# compute discriminator loss for real mu maps
outputs_d_real = self.discriminator(inputs_d_real)
loss_d_real = self.criterion_adv(outputs_d_real, labels_real)
# update discriminator
loss_d = 0.5 * (loss_d_fake + loss_d_real)
loss_d.backward() # compute gradients
self.optim_d.step()
# ===========================================================================
# ======================= Generator =========================================
outputs_d_fake = self.discriminator(inputs_d_fake) # this time no detach
loss_g_adv = self.criterion_adv(outputs_d_fake, labels_real)
loss_g_dist = self.criterion_dist(mu_maps_fake, mu_maps_real)
loss_g = (
self.weight_criterion_adv * loss_g_adv
+ self.weight_criterion_dist * loss_g_dist
)
loss_g.backward()
self.optim_g.step()
# ===========================================================================
def _store_snapshot(self, epoch):
prefix = f"{epoch:0{len(str(self.epochs))}d}"
self.store_snapshot(prefix)
return loss_g.item()
def store_snapshot(self, prefix: str):
snapshot_file_d = os.path.join(self.snapshot_dir, f"{prefix}_discriminator.pth")
snapshot_file_g = os.path.join(self.snapshot_dir, f"{prefix}_generator.pth")
self.logger.debug(f"Store snapshots at {snapshot_file_d} and {snapshot_file_g}")
torch.save(self.params_d.model.state_dict(), snapshot_file_d)
torch.save(self.params_g.model.state_dict(), snapshot_file_g)
def _eval_batch(self, recons: torch.Tensor, mu_maps: torch.Tensor) -> float:
mu_maps_fake = self.generator(recons)
loss = torch.nn.functional.l1_loss(mu_maps_fake, mu_maps)
return loss.item()
if __name__ == "__main__":
import argparse
import os
import random
import sys
......@@ -198,7 +170,7 @@ if __name__ == "__main__":
from mu_map.models.discriminator import Discriminator, PatchDiscriminator
parser = argparse.ArgumentParser(
description="Train a UNet model to predict μ-maps from reconstructed scatter images",
description="Train a UNet model to predict μ-maps from reconstructed images",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
......@@ -334,11 +306,6 @@ if __name__ == "__main__":
default=10,
help="frequency in epochs at which snapshots are stored",
)
parser.add_argument(
"--generator_weights",
type=str,
help="use pre-trained weights for the generator",
)
# Logging Args
add_logging_args(parser, defaults={"--logfile": "train.log"})
......@@ -361,10 +328,11 @@ if __name__ == "__main__":
args.logfile = os.path.join(args.output_dir, args.logfile)
device = torch.device(args.device)
logger = get_logger_by_args(args)
logger.info(args)
device = torch.device(args.device)
args.seed = args.seed if args.seed is not None else random.randint(0, 2**32 - 1)
logger.info(f"Seed: {args.seed}")
random.seed(args.seed)
......@@ -382,29 +350,19 @@ if __name__ == "__main__":
[transform_normalization, PadCropTranform(dim=3, size=32)]
)
data_loaders = {}
for split in ["train", "validation"]:
dataset = MuMapPatchDataset(
args.dataset_dir,
patches_per_image=args.number_of_patches,
patch_size=args.patch_size,
patch_offset=args.patch_offset,
shuffle=not args.no_shuffle,
split_name=split,
transform_normalization=transform_normalization,
scatter_correction=args.scatter_correction,
logger=logger,
)
data_loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=1,
)
data_loaders[split] = data_loader
dataset = MuMapPatchDataset(
args.dataset_dir,
patches_per_image=args.number_of_patches,
patch_size=args.patch_size,
patch_offset=args.patch_offset,
shuffle=not args.no_shuffle,
transform_normalization=transform_normalization,
scatter_correction=args.scatter_correction,
logger=logger,
)
discriminator = Discriminator(in_channels=2, input_size=args.patch_size)
# discriminator = PatchDiscriminator(in_channels=2)
discriminator = discriminator.to(device)
optimizer = torch.optim.Adam(
discriminator.parameters(), lr=args.lr, betas=(0.5, 0.999)
......@@ -416,17 +374,12 @@ if __name__ == "__main__":
if args.decay_lr
else None
)
params_d = TrainingParams(
params_d = DiscriminatorParams(
model=discriminator, optimizer=optimizer, lr_scheduler=lr_scheduler
)
generator = UNet(in_channels=1, features=args.features)
generator = generator.to(device)
if args.generator_weights:
logger.debug(f"Load generator weights from {args.generator_weights}")
generator.load_state_dict(
torch.load(args.generator_weights, map_location=device)
)
optimizer = torch.optim.Adam(generator.parameters(), lr=args.lr, betas=(0.5, 0.999))
lr_scheduler = (
torch.optim.lr_scheduler.StepLR(
......@@ -435,7 +388,7 @@ if __name__ == "__main__":
if args.decay_lr
else None
)
params_g = TrainingParams(
params_g = GeneratorParams(
model=generator, optimizer=optimizer, lr_scheduler=lr_scheduler
)
......@@ -443,16 +396,17 @@ if __name__ == "__main__":
logger.debug(f"Use distance criterion: {dist_criterion}")
training = cGANTraining(
data_loaders=data_loaders,
epochs=args.epochs,
dataset=dataset,
batch_size=args.batch_size,
device=device,
snapshot_dir=args.snapshot_dir,
snapshot_epoch=args.snapshot_epoch,
logger=logger,
params_generator=params_g,
params_discriminator=params_d,
loss_func_dist=dist_criterion,
weight_criterion_dist=args.dist_loss_weight,
weight_criterion_adv=args.adv_loss_weight,
logger=logger,
)
training.run()
mu_map/training/cgan2.py deleted 100644 → 0
+ 0
− 412
View file @ 7fd7c7b2
from logging import Logger
from typing import Optional
import torch
from mu_map.dataset.default import MuMapDataset
from mu_map.training.lib import TrainingParams, AbstractTraining
from mu_map.training.loss import WeightedLoss
# Establish convention for real and fake labels during training
LABEL_REAL = 1.0
LABEL_FAKE = 0.0
class DiscriminatorParams(TrainingParams):
"""
Wrap training parameters to always carry the name 'Discriminator'.
"""
def __init__(
self,
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
):
super().__init__(
name="Discriminator",
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
)
class GeneratorParams(TrainingParams):
"""
Wrap training parameters to always carry the name 'Generator'.
"""
def __init__(
self,
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
):
super().__init__(
name="Generator",
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
)
class cGANTraining(AbstractTraining):
"""
Implementation of a conditional generative adversarial network training.
"""
def __init__(
self,
epochs: int,
dataset: MuMapDataset,
batch_size: int,
device: torch.device,
snapshot_dir: str,
snapshot_epoch: int,
params_generator: GeneratorParams,
params_discriminator: DiscriminatorParams,
loss_func_dist: WeightedLoss,
weight_criterion_dist: float,
weight_criterion_adv: float,
logger: Optional[Logger] = None,
):
"""
:param params_generator: training parameters containing a model an according optimizer and optionally a learning rate scheduler for the generator
:param params_discriminator: training parameters containing a model an according optimizer and optionally a learning rate scheduler for the discriminator
:param loss_func_dist: distance loss function for the generator
:param weight_criterion_dist: weight of the distance loss when training the generator
:param weight_criterion_adv: weight of the adversarial loss when training the generator
"""
super().__init__(
epochs, dataset, batch_size, device, snapshot_dir, snapshot_epoch, logger
)
self.training_params.append(params_generator)
self.training_params.append(params_discriminator)
self.generator = params_generator.model
self.discriminator = params_discriminator.model
self.optim_g = params_generator.optimizer
self.optim_d = params_discriminator.optimizer
self.weight_criterion_dist = weight_criterion_dist
self.weight_criterion_adv = weight_criterion_adv
self.criterion_adv = torch.nn.MSELoss(reduction="mean")
self.criterion_dist = loss_func_dist
def _after_train_batch(self):
"""
Overwrite calling step on all optimizers as this needs to be done
separately for the generator and discriminator during the training of
a batch.
"""
pass
def _train_batch(self, recons: torch.Tensor, mu_maps: torch.Tensor) -> float:
mu_maps_real = mu_maps # rename real mu maps for clarification
# compute fake mu maps with generator
mu_maps_fake = self.generator(recons)
# prepare inputs for the discriminator
inputs_d_fake = torch.cat((recons, mu_maps_fake), dim=1)
inputs_d_real = torch.cat((recons, mu_maps_real), dim=1)
# prepare labels/targets for the discriminator
labels_fake = torch.full(self.discriminator.get_output_shape(inputs_d_fake.shape), LABEL_FAKE, device=self.device)
labels_real = torch.full(self.discriminator.get_output_shape(inputs_d_real.shape), LABEL_REAL, device=self.device)
# ======================= Discriminator =====================================
# compute discriminator loss for fake mu maps
# detach is called so that gradients are not computed for the generator
outputs_d_fake = self.discriminator(inputs_d_fake.detach())
loss_d_fake = self.criterion_adv(outputs_d_fake, labels_fake)
# compute discriminator loss for real mu maps
outputs_d_real = self.discriminator(inputs_d_real)
loss_d_real = self.criterion_adv(outputs_d_real, labels_real)
# update discriminator
loss_d = 0.5 * (loss_d_fake + loss_d_real)
loss_d.backward() # compute gradients
self.optim_d.step()
# ===========================================================================
# ======================= Generator =========================================
outputs_d_fake = self.discriminator(inputs_d_fake) # this time no detach
loss_g_adv = self.criterion_adv(outputs_d_fake, labels_real)
loss_g_dist = self.criterion_dist(mu_maps_fake, mu_maps_real)
loss_g = (
self.weight_criterion_adv * loss_g_adv
+ self.weight_criterion_dist * loss_g_dist
)
loss_g.backward()
self.optim_g.step()
# ===========================================================================
return loss_g.item()
def _eval_batch(self, recons: torch.Tensor, mu_maps: torch.Tensor) -> float:
mu_maps_fake = self.generator(recons)
loss = torch.nn.functional.l1_loss(mu_maps_fake, mu_maps)
return loss.item()
if __name__ == "__main__":
import argparse
import os
import random
import sys
import numpy as np
from mu_map.dataset.patches import MuMapPatchDataset
from mu_map.dataset.normalization import (
MeanNormTransform,
MaxNormTransform,
GaussianNormTransform,
)
from mu_map.dataset.transform import PadCropTranform, SequenceTransform
from mu_map.logging import add_logging_args, get_logger_by_args
from mu_map.models.unet import UNet
from mu_map.models.discriminator import Discriminator, PatchDiscriminator
parser = argparse.ArgumentParser(
description="Train a UNet model to predict μ-maps from reconstructed images",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
# Model Args
parser.add_argument(
"--features",
type=int,
nargs="+",
default=[64, 128, 256, 512],
help="number of features in the layers of the UNet structure",
)
# Dataset Args
parser.add_argument(
"--dataset_dir",
type=str,
default="data/second/",
help="the directory where the dataset for training is found",
)
parser.add_argument(
"--input_norm",
type=str,
choices=["none", "mean", "max", "gaussian"],
default="mean",
help="type of normalization applied to the reconstructions",
)
parser.add_argument(
"--patch_size",
type=int,
default=32,
help="the size of patches extracted for each reconstruction",
)
parser.add_argument(
"--patch_offset",
type=int,
default=20,
help="offset to ignore the border of the image",
)
parser.add_argument(
"--number_of_patches",
type=int,
default=100,
help="number of patches extracted for each image",
)
parser.add_argument(
"--no_shuffle",
action="store_true",
help="do not shuffle patches in the dataset",
)
parser.add_argument(
"--scatter_correction",
action="store_true",
help="use the scatter corrected reconstructions in the dataset",
)
# Training Args
parser.add_argument(
"--seed",
type=int,
help="seed used for random number generation",
)
parser.add_argument(
"--batch_size",
type=int,
default=64,
help="the batch size used for training",
)
parser.add_argument(
"--output_dir",
type=str,
default="train_data",
help="directory in which results (snapshots and logs) of this training are saved",
)
parser.add_argument(
"--epochs",
type=int,
default=100,
help="the number of epochs for which the model is trained",
)
parser.add_argument(
"--device",
type=str,
default="cuda:0" if torch.cuda.is_available() else "cpu",
help="the device (cpu or gpu) with which the training is performed",
)
parser.add_argument(
"--dist_loss_func",
type=str,
default="l1",
help="define the loss function used as the distance loss of the generator , e.g. 0.75*l2+0.25*gdl",
)
parser.add_argument(
"--dist_loss_weight",
type=float,
default=100.0,
help="weight for the distance loss of the generator",
)
parser.add_argument(
"--adv_loss_weight",
type=float,
default=1.0,
help="weight for the Adversarial-Loss of the generator",
)
parser.add_argument(
"--lr", type=float, default=0.001, help="the initial learning rate for training"
)
parser.add_argument(
"--decay_lr",
action="store_true",
help="decay the learning rate",
)
parser.add_argument(
"--lr_decay_factor",
type=float,
default=0.99,
help="decay factor for the learning rate",
)
parser.add_argument(
"--lr_decay_epoch",
type=int,
default=1,
help="frequency in epochs at which the learning rate is decayed",
)
parser.add_argument(
"--snapshot_dir",
type=str,
default="snapshots",
help="directory under --output_dir where snapshots are stored",
)
parser.add_argument(
"--snapshot_epoch",
type=int,
default=10,
help="frequency in epochs at which snapshots are stored",
)
# Logging Args
add_logging_args(parser, defaults={"--logfile": "train.log"})
args = parser.parse_args()
if not os.path.exists(args.output_dir):
os.mkdir(args.output_dir)
args.snapshot_dir = os.path.join(args.output_dir, args.snapshot_dir)
if not os.path.exists(args.snapshot_dir):
os.mkdir(args.snapshot_dir)
else:
if len(os.listdir(args.snapshot_dir)) > 0:
print(
f"ATTENTION: Snapshot directory [{args.snapshot_dir}] already exists and is not empty!"
)
print(f" Exit so that data is not accidentally overwritten!")
exit(1)
args.logfile = os.path.join(args.output_dir, args.logfile)
logger = get_logger_by_args(args)
logger.info(args)
device = torch.device(args.device)
args.seed = args.seed if args.seed is not None else random.randint(0, 2**32 - 1)
logger.info(f"Seed: {args.seed}")
random.seed(args.seed)
torch.manual_seed(args.seed)
np.random.seed(args.seed)
transform_normalization = None
if args.input_norm == "mean":
transform_normalization = MeanNormTransform()
elif args.input_norm == "max":
transform_normalization = MaxNormTransform()
elif args.input_norm == "gaussian":
transform_normalization = GaussianNormTransform()
transform_normalization = SequenceTransform(
[transform_normalization, PadCropTranform(dim=3, size=32)]
)
dataset = MuMapPatchDataset(
args.dataset_dir,
patches_per_image=args.number_of_patches,
patch_size=args.patch_size,
patch_offset=args.patch_offset,
shuffle=not args.no_shuffle,
transform_normalization=transform_normalization,
scatter_correction=args.scatter_correction,
logger=logger,
)
discriminator = Discriminator(in_channels=2, input_size=args.patch_size)
# discriminator = PatchDiscriminator(in_channels=2)
discriminator = discriminator.to(device)
optimizer = torch.optim.Adam(
discriminator.parameters(), lr=args.lr, betas=(0.5, 0.999)
)
lr_scheduler = (
torch.optim.lr_scheduler.StepLR(
optimizer, step_size=args.lr_decay_epoch, gamma=args.lr_decay_factor
)
if args.decay_lr
else None
)
params_d = DiscriminatorParams(
model=discriminator, optimizer=optimizer, lr_scheduler=lr_scheduler
)
generator = UNet(in_channels=1, features=args.features)
generator = generator.to(device)
optimizer = torch.optim.Adam(generator.parameters(), lr=args.lr, betas=(0.5, 0.999))
lr_scheduler = (
torch.optim.lr_scheduler.StepLR(
optimizer, step_size=args.lr_decay_factor, gamma=args.lr_decay_factor
)
if args.decay_lr
else None
)
params_g = GeneratorParams(
model=generator, optimizer=optimizer, lr_scheduler=lr_scheduler
)
dist_criterion = WeightedLoss.from_str(args.dist_loss_func)
logger.debug(f"Use distance criterion: {dist_criterion}")
training = cGANTraining(
epochs=args.epochs,
dataset=dataset,
batch_size=args.batch_size,
device=device,
snapshot_dir=args.snapshot_dir,
snapshot_epoch=args.snapshot_epoch,
params_generator=params_g,
params_discriminator=params_d,
loss_func_dist=dist_criterion,
weight_criterion_dist=args.dist_loss_weight,
weight_criterion_adv=args.adv_loss_weight,
logger=logger,
)
training.run()
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment