def prepare_generator(writer, dataset, device): generator = modules.Generator(in_channels=3) optimizer = torch.optim.SGD(generator.parameters(), lr=0.0001) step = (steps.Generator( tt.loss.SmoothBinaryCrossEntropy(alpha=0.1), device)**tt.Select( loss=0)**tt.pytorch.ZeroGrad(optimizer)**tt.pytorch.Backward()** tt.pytorch.Optimize(optimizer)**tt.pytorch.Detach()) step**tt.Select(generated_images=2)**tt.callbacks.tensorboard.Images( writer, "Generator/Images")**operations.AddFakeImages(dataset) return generator, step
def accuracy_pipeline(iteration, writer, name: str): return (iteration > tt.Select(predictions=1, labels=2) > tt.metrics.classification.multiclass.Accuracy() > tt.accumulators.Mean() > tt.Split( tt.callbacks.Log(f"{name} Accuracy"), tt.callbacks.tensorboard.Scalar(writer, f"{name}/Accuracy"), ))
def prepare_discriminator(device): discriminator = modules.Discriminator(in_channels=3) optimizer = torch.optim.SGD(discriminator.parameters(), lr=0.0004) step = (steps.Discriminator( tt.loss.SmoothBinaryCrossEntropy(alpha=0.1), device)**tt.Select( loss=0)**tt.pytorch.ZeroGrad(optimizer)**tt.pytorch.Backward()** tt.pytorch.Optimize(optimizer)**tt.pytorch.Detach()) return discriminator, step
def test_step(): module = torch.nn.Linear(10, 1) optimizer = torch.optim.Adam(module.parameters()) step = (Step(torch.nn.MSELoss())**tt.Select( loss=0)**tt.pytorch.ZeroGrad(optimizer)**tt.pytorch.Backward( accumulate=2)**tt.pytorch.Optimize( optimizer, accumulate=2)**tt.pytorch.Detach()**tt.callbacks.Log("Loss")) step = Step(torch.nn.MSELoss())**tt.Select(loss=0) step**tt.Select( predictions=1, target=2)**tt.metrics.regression.MaxError()**tt.callbacks.Log( name="Max Error") for _ in range(5): step(module, (torch.randn(8, 10), torch.randn(8)))
def prepare_iteration( writer, generator, generator_step, discriminator, discriminator_step, noise_dataset, cifar10_with_fake, ): iteration = tt.iterations.MultiIteration( steps=(generator_step, discriminator_step), modules=((generator, discriminator), discriminator), datas=(noise_dataset, cifar10_with_fake), intervals=(4, 1), log="INFO", ) iteration**tt.Select( loss=0)**tt.pytorch.Detach()**tt.device.CPU()**tt.Except( tt.accumulators.Mean(), 4)**tt.Split( tt.callbacks.tensorboard.Scalar(writer, "Generator/Loss"), tt.callbacks.Logger(name="Generator Mean"), tt.callbacks.Save(generator, "generator.pt", comparator=operator.lt), ) iteration**tt.Select( loss=0)**tt.pytorch.Detach()**tt.device.CPU()**tt.Except( tt.accumulators.Mean(), begin=0, end=4)**tt.Split( tt.callbacks.tensorboard.Scalar(writer, "Discriminator/Loss"), tt.callbacks.Logger(name="Generator Mean"), tt.callbacks.Save(discriminator, "generator.pt", comparator=operator.lt), ) return iteration
def train(optimizer, criterion, device): return (Step( criterion, gradient=True, device=device)**tt.Select(loss=0)**tt.pytorch.ZeroGrad(optimizer)** tt.pytorch.Backward()**tt.pytorch.Optimize(optimizer))
def loss_pipeline(iteration, writer, name: str): return (iteration > tt.Select(loss=0) > tt.accumulators.Mean() > tt.Split( tt.callbacks.Log(f"{name} Loss"), tt.callbacks.tensorboard.Scalar(writer, f"{name}/Loss", log="INFO"), ))