def train_net(Gen, Discr): G = Gen(in_noise=128, out_ch=3) G_polyak = copy.deepcopy(G).eval() D = Discr() print(G) print(D) def G_fun(batch): z = torch.randn(BS, 128, device=device) fake = G(z) preds = D(fake * 2 - 1).squeeze() loss = gan_loss.generated(preds) loss.backward() return {'loss': loss.item(), 'imgs': fake.detach()} def G_polyak_fun(batch): z = torch.randn(BS, 128, device=device) fake = G_polyak(z) return {'imgs': fake.detach()} def D_fun(batch): z = torch.randn(BS, 128, device=device) fake = G(z) fake_loss = gan_loss.fake(D(fake * 2 - 1)) fake_loss.backward() x = batch[0] real_loss = gan_loss.real(D(x * 2 - 1)) real_loss.backward() loss = real_loss.item() + fake_loss.item() return { 'loss': loss, 'real_loss': real_loss.item(), 'fake_loss': fake_loss.item() } loop = GANRecipe(G, D, G_fun, D_fun, G_polyak_fun, dl, log_every=100).to(device) loop.register('polyak', G_polyak) loop.G_loop.callbacks.add_callbacks([ tcb.Optimizer( tch.optim.RAdamW(G.parameters(), lr=1e-4, betas=(0., 0.99))), tcb.Polyak(G, G_polyak), ]) loop.register('G_polyak', G_polyak) loop.callbacks.add_callbacks([ tcb.Log('batch.0', 'x'), tcb.WindowedMetricAvg('real_loss'), tcb.WindowedMetricAvg('fake_loss'), tcb.Optimizer( tch.optim.RAdamW(D.parameters(), lr=4e-4, betas=(0., 0.99))), ]) loop.test_loop.callbacks.add_callbacks([ tcb.Log('imgs', 'polyak_imgs'), tcb.VisdomLogger('main', prefix='test') ]) loop.to(device).run(100)
def fit(self, iters, content_img, style_img, style_ratio, content_layers=None): """ Run the recipe Args: n_iters (int): number of iterations to run content (PIL.Image): content image style (PIL.Image): style image ratio (float): weight of style loss content_layers (list of str): layers on which to reconstruct content """ self.loss.to(self.device) self.loss.set_style(pil2t(style_img).to(self.device), style_ratio) self.loss.set_content(pil2t(content_img).to(self.device), content_layers) canvas = ParameterizedImg(3, content_img.height, content_img.width, init_sd=0.00) self.opt = tch.optim.RAdamW(canvas.parameters(), 1e-2, (0.7, 0.7), eps=0.00001, weight_decay=0) def forward(_): self.opt.zero_grad() img = canvas() loss, losses = self.loss(img) loss.backward() return { 'loss': loss, 'content_loss': losses['content_loss'], 'style_loss': losses['style_loss'], 'img': img } loop = Recipe(forward, range(iters)) loop.register('canvas', canvas) loop.register('model', self) loop.callbacks.add_callbacks([ tcb.Counter(), tcb.WindowedMetricAvg('loss'), tcb.WindowedMetricAvg('content_loss'), tcb.WindowedMetricAvg('style_loss'), tcb.Log('img', 'img'), tcb.VisdomLogger(visdom_env=self.visdom_env, log_every=10), tcb.StdoutLogger(log_every=10), tcb.Optimizer(self.opt, log_lr=True), tcb.LRSched(torch.optim.lr_scheduler.ReduceLROnPlateau(self.opt, threshold=0.001, cooldown=500), step_each_batch=True) ]) loop.to(self.device) loop.run(1) return canvas.render().cpu()
def make_loop(hourglass, body, display, num_iter, lr): loop = TrainAndCall(hourglass, body, display, range(num_iter), test_every=50, checkpoint=None) opt = tch.optim.RAdamW(hourglass.parameters(), lr=lr) loop.callbacks.add_callbacks([ tcb.WindowedMetricAvg('loss'), tcb.Optimizer(opt, clip_grad_norm=0.5, log_lr=True), ]) loop.test_loop.callbacks.add_callbacks([ tcb.Log('recon', 'img'), tcb.Log('orig', 'orig'), tcb.Log('loss', 'loss'), ]) return loop
def fit(self, n_iters, neuron): """ Run the recipe Args: n_iters (int): number of iterations to run neuron (int): the feature map to maximize Returns: the optimized image """ canvas = ParameterizedImg(3, self.input_size + 10, self.input_size + 10) def forward(_): cim = canvas() rnd = random.randint(0, cim.shape[2] // 10) im = cim[:, :, rnd:, rnd:] im = torch.nn.functional.interpolate(im, size=(self.input_size, self.input_size), mode='bilinear') _, acts = self.model(self.norm(im), detach=False) fmap = acts[self.layer] loss = -fmap[0][neuron].sum() loss.backward() return {'loss': loss, 'img': cim} loop = Recipe(forward, range(n_iters)) loop.register('canvas', canvas) loop.register('model', self) loop.callbacks.add_callbacks([ tcb.Counter(), tcb.Log('loss', 'loss'), tcb.Log('img', 'img'), tcb.Optimizer(DeepDreamOptim(canvas.parameters(), lr=self.lr)), tcb.VisdomLogger(visdom_env=self.visdom_env, log_every=10), tcb.StdoutLogger(log_every=10) ]) loop.to(self.device) loop.run(1) return canvas.render().cpu()
def test_trainandcall(): model = nn.Linear(10, 2) def train_step(batch): x, y = batch out = model(x) loss = torch.nn.functional.cross_entropy(out, y) loss.backward() return {'loss': loss} def after_train(): print('Yup.') return {} trainloader = DataLoader(FakeData(), 4, shuffle=True) trainer = TrainAndCall(model, train_step, after_train, trainloader) trainer.callbacks.add_callbacks( [tcb.Optimizer(torch.optim.Adam(model.parameters(), lr=1e-3))]) trainer.run(1)
def fit(self, ref, iters, lr=3e-4, device='cpu', visdom_env='deepdream'): """ Args: lr (float, optional): the learning rate visdom_env (str or None): the name of the visdom env to use, or None to disable Visdom """ ref_tensor = TF.ToTensor()(ref).unsqueeze(0) canvas = ParameterizedImg(1, 3, ref_tensor.shape[2], ref_tensor.shape[3], init_img=ref_tensor, space='spectral', colors='uncorr') def forward(_): img = canvas() rnd = random.randint(0, 10) loss = self.loss(self.norm(img[:, :, rnd:, rnd:])) loss.backward() return {'loss': loss, 'img': img} loop = Recipe(forward, range(iters)) loop.register('model', self) loop.register('canvas', canvas) loop.callbacks.add_callbacks([ tcb.Counter(), tcb.Log('loss', 'loss'), tcb.Log('img', 'img'), tcb.Optimizer(DeepDreamOptim(canvas.parameters(), lr=lr)), tcb.VisdomLogger(visdom_env=visdom_env, log_every=10), tcb.StdoutLogger(log_every=10) ]) loop.to(device) loop.run(1) return canvas.render().cpu()
def CrossEntropyClassification(model, train_loader, test_loader, classes, lr=3e-3, beta1=0.9, wd=1e-2, visdom_env='main', test_every=1000, log_every=100): """ A Classification recipe with a default froward training / testing pass using cross entropy, and extended with RAdamW and ReduceLROnPlateau. Args: model (nn.Module): a model learnable with cross entropy train_loader (DataLoader): Training set dataloader test_loader (DataLoader): Testing set dataloader classes (list of str): classes name, in order lr (float): the learning rate beta1 (float): RAdamW's beta1 wd (float): weight decay visdom_env (str): name of the visdom environment to use, or None for not using Visdom (default: None) test_every (int): testing frequency, in number of iterations (default: 1000) log_every (int): logging frequency, in number of iterations (default: 1000) """ def train_step(batch): x, y = batch pred = model(x) loss = torch.nn.functional.cross_entropy(pred, y) loss.backward() return {'loss': loss, 'pred': pred} def validation_step(batch): x, y = batch pred = model(x) loss = torch.nn.functional.cross_entropy(pred, y) return {'loss': loss, 'pred': pred} loop = Classification(model, train_step, validation_step, train_loader, test_loader, classes, visdom_env=visdom_env, test_every=test_every, log_every=log_every) opt = RAdamW(model.parameters(), lr=lr, betas=(beta1, 0.999), weight_decay=wd) loop.callbacks.add_callbacks([ tcb.Optimizer(opt, log_lr=True), tcb.LRSched(torch.optim.lr_scheduler.ReduceLROnPlateau(opt)) ]) return loop
def fit(self, iters, content_img, style_img, style_ratio, content_layers=None): """ Run the recipe Args: n_iters (int): number of iterations to run content (PIL.Image): content image style (PIL.Image): style image ratio (float): weight of style loss content_layers (list of str): layers on which to reconstruct content """ self.loss.to(self.device) self.loss.set_style(pil2t(style_img).to(self.device), style_ratio) self.loss.set_content( pil2t(content_img).to(self.device), content_layers) self.loss2.to(self.device) self.loss2.set_style( torch.nn.functional.interpolate(pil2t(style_img)[None], scale_factor=0.5, mode='bilinear')[0].to( self.device), style_ratio) self.loss2.set_content( torch.nn.functional.interpolate(pil2t(content_img)[None], scale_factor=0.5, mode='bilinear')[0].to( self.device), content_layers) canvas = ParameterizedImg(3, content_img.height, content_img.width, init_img=pil2t(content_img)) self.opt = tch.optim.RAdamW(canvas.parameters(), 3e-2) def forward(_): img = canvas() loss, losses = self.loss(img) loss.backward() loss, losses = self.loss2( torch.nn.functional.interpolate(canvas(), scale_factor=0.5, mode='bilinear')) loss.backward() return { 'loss': loss, 'content_loss': losses['content_loss'], 'style_loss': losses['style_loss'], 'img': img } loop = Recipe(forward, range(iters)) loop.register('canvas', canvas) loop.register('model', self) loop.callbacks.add_callbacks([ tcb.Counter(), tcb.WindowedMetricAvg('loss'), tcb.WindowedMetricAvg('content_loss'), tcb.WindowedMetricAvg('style_loss'), tcb.Log('img', 'img'), tcb.VisdomLogger(visdom_env=self.visdom_env, log_every=10), tcb.StdoutLogger(log_every=10), tcb.Optimizer(self.opt, log_lr=True), ]) loop.to(self.device) loop.run(1) return canvas.render().cpu()
def MixupClassification(model, train_loader, test_loader, classes, *, lr=3e-3, beta1=0.9, wd=1e-2, visdom_env='main', test_every=1000, log_every=100): """ A Classification recipe with a default froward training / testing pass using cross entropy and mixup, and extended with RAdamW and ReduceLROnPlateau. Args: model (nn.Module): a model learnable with cross entropy train_loader (DataLoader): Training set dataloader. Must have soft targets. Should be a DataLoader loading a MixupDataset or compatible. test_loader (DataLoader): Testing set dataloader. Dataset must have categorical targets. classes (list of str): classes name, in order lr (float): the learning rate beta1 (float): RAdamW's beta1 wd (float): weight decay visdom_env (str): name of the visdom environment to use, or None for not using Visdom (default: None) test_every (int): testing frequency, in number of iterations (default: 1000) log_every (int): logging frequency, in number of iterations (default: 1000) """ from torchelie.loss import continuous_cross_entropy def train_step(batch): x, y = batch pred = model(x) loss = continuous_cross_entropy(pred, y) loss.backward() return {'loss': loss} def validation_step(batch): x, y = batch pred = model(x) loss = torch.nn.functional.cross_entropy(pred, y) return {'loss': loss, 'pred': pred} loop = TrainAndTest(model, train_step, validation_step, train_loader, test_loader, visdom_env=visdom_env, test_every=test_every, log_every=log_every) loop.callbacks.add_callbacks([ tcb.WindowedMetricAvg('loss'), ]) loop.register('classes', classes) loop.test_loop.callbacks.add_callbacks([ tcb.AccAvg(post_each_batch=False), tcb.WindowedMetricAvg('loss', False), ]) if visdom_env is not None: loop.callbacks.add_epilogues( [tcb.ImageGradientVis(), tcb.MetricsTable()]) if len(classes) <= 25: loop.test_loop.callbacks.add_callbacks([ tcb.ConfusionMatrix(classes), ]) loop.test_loop.callbacks.add_callbacks([ tcb.ClassificationInspector(30, classes, False), tcb.MetricsTable(False) ]) opt = RAdamW(model.parameters(), lr=lr, betas=(beta1, 0.999), weight_decay=wd) loop.callbacks.add_callbacks([ tcb.Optimizer(opt, log_lr=True), tcb.LRSched(torch.optim.lr_scheduler.ReduceLROnPlateau(opt)) ]) return loop
def CrossEntropyClassification(model, train_loader, test_loader, classes, lr=3e-3, beta1=0.9, wd=1e-2, visdom_env='main', test_every=1000, log_every=100): """ Extends Classification with default cross entropy forward passes. Also adds RAdamW and ReduceLROnPlateau. Inherited training callbacks: - AccAvg for displaying accuracy - WindowedMetricAvg for displaying loss - ConfusionMatrix if len(classes) <= 25 - ClassificationInspector - MetricsTable - ImageGradientVis - Counter for counting iterations, connected to the testing loop as well - VisdomLogger - StdoutLogger Training callbacks: - Optimizer with RAdamW - LRSched with ReduceLROnPlateau Testing: Testing loop is in :code:`.test_loop`. Inherited testing callbacks: - AccAvg - WindowedMetricAvg - ConfusionMatrix if :code:`len(classes) <= 25` - ClassificationInspector - MetricsTable - VisdomLogger - StdoutLogger - Checkpoint saving the best testing loss Args: model (nn.Module): a model learnable with cross entropy train_loader (DataLoader): Training set dataloader test_loader (DataLoader): Testing set dataloader classes (list of str): classes name, in order lr (float): the learning rate beta1 (float): RAdamW's beta1 wd (float): weight decay visdom_env (str): name of the visdom environment to use, or None for not using Visdom (default: None) test_every (int): testing frequency, in number of iterations (default: 1000) log_every (int): logging frequency, in number of iterations (default: 1000) """ def train_step(batch): x, y = batch pred = model(x) loss = torch.nn.functional.cross_entropy(pred, y) loss.backward() return {'loss': loss, 'pred': pred} def validation_step(batch): x, y = batch pred = model(x) loss = torch.nn.functional.cross_entropy(pred, y) return {'loss': loss, 'pred': pred} loop = Classification(model, train_step, validation_step, train_loader, test_loader, classes, visdom_env=visdom_env, test_every=test_every, log_every=log_every) opt = RAdamW(model.parameters(), lr=lr, betas=(beta1, 0.999), weight_decay=wd) loop.callbacks.add_callbacks([ tcb.Optimizer(opt, log_lr=True), tcb.LRSched(torch.optim.lr_scheduler.ReduceLROnPlateau(opt)) ]) return loop