from torch.utils.tensorboard import SummaryWriter from dataset.lazy_loader import LazyLoader, W300DatasetLoader, CelebaWithKeyPoints, Celeba from dataset.probmeasure import ProbabilityMeasureFabric, ProbabilityMeasure, UniformMeasure2DFactory from gan.gan_model import cont_style_munit_enc, CondGen3, CondDisc3, CondStyleGanModel from metrics.writers import ItersCounter, send_images_to_tensorboard from nn.munit.enc_dec import MunitEncoder, StyleEncoder from modules.hg import hg2, final_preds_untransformed, hg8, hg4, HG_softmax2020 from gan.loss_base import Loss from transforms_utils.transforms import MeasureToMask, ToNumpy, ToTensor, MaskToMeasure, NumpyBatch, MeasureToKeyPoints, \ ResizeMask, NormalizeMask device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(device) torch.cuda.set_device(device) counter = ItersCounter() writer = SummaryWriter( f"/home/ibespalov/pomoika/KKKKKstylegan{int(time.time())}") def writable(name: str, f: Callable[[Any], Loss]): counter.active[name] = True def decorated(*args, **kwargs) -> Loss: loss = f(*args, **kwargs) it = counter.get_iter(name) if it % 10 == 0: writer.add_scalar(name, loss.item(), it) return loss return decorated
# nn.MSELoss()(pred[0].coord, t_xy) * (0.0005 * coef) + # nn.L1Loss()(pred[0].coord, t_xy) * (0.0005 * coef) + # (pred[1] - target).abs().mean() * (0.3 * coef) # ) # def rt_loss(pred: Tuple[UniformMeasure2D01, Tensor], target: Tensor): # # with torch.no_grad(): # t_xy, _ = heatmap_to_measure(target) # # return Loss( # nn.L1Loss()(pred[0].coord, t_xy) # ) counter = ItersCounter() writer = SummaryWriter(f"{Paths.default.board()}/stylegan{int(time.time())}") print(f"{Paths.default.board()}/stylegan{int(time.time())}") l1_loss = nn.L1Loss() def L1(name: Optional[str], writer: SummaryWriter = writer) -> Callable[[Tensor, Tensor], Loss]: if name: counter.active[name] = True def compute(t1: Tensor, t2: Tensor): loss = l1_loss(t1, t2) if name: if counter.get_iter(name) % 10 == 0: writer.add_scalar(name, loss, counter.get_iter(name)) return Loss(loss)
from torch import optim from torch import nn, Tensor from torch.utils.tensorboard import SummaryWriter from dataset.lazy_loader import LazyLoader, W300DatasetLoader, Celeba from dataset.probmeasure import ProbabilityMeasureFabric, ProbabilityMeasure, UniformMeasure2DFactory from metrics.writers import ItersCounter, send_images_to_tensorboard from modules.hg import HG_softmax2020 from gan.loss_base import Loss from matplotlib import pyplot as plt device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(device) torch.cuda.set_device(device) counter = ItersCounter() writer = SummaryWriter(f"{Paths.default.board()}/w300{int(time.time())}") print(f"{Paths.default.board()}/w300{int(time.time())}") def writable(name: str, f: Callable[[Any], Loss]): counter.active[name] = True def decorated(*args, **kwargs) -> Loss: loss = f(*args, **kwargs) writer.add_scalar(name, loss.item(), counter.get_iter(name)) return loss return decorated