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