code: # coding:utf8
  
  import torch as t
  import torchvision as tv
  import torchnet as tnt
  
  from torch.utils import data
  from transformer_net import TransformerNet
  import utils
  from PackedVGG import Vgg16
  from torch.nn import functional as F
  import tqdm
  import os
  import ipdb
  # from WCT2_train import WCT2
  # import model
  
  from LapSobGaus_train import Lap_Sob_Gaus
  import net
  
  import Ovodus_Laplace_model
  import utils_
  from WCT2_train import train_transform
  from tensorboardX import SummaryWriter
  
  from pathlib import Path
  
  from torchvision.utils import save_image
  
  mean = [0.485, 0.456, 0.406]
  std = [0.229, 0.224, 0.225]
  
  IMAGENET_MEAN = [0.485, 0.456, 0.406]
  IMAGENET_STD = [0.229, 0.224, 0.225]
  
  
  class Config(object):
      # General Args
      use_gpu = True
      model_path = None  # pretrain model path (for resume training or test)
  
      # Train Args
      image_size = 448  # image crop_size for training
      batch_size = 2
      data_root = r'F:\DataSets\train2017'  # 'data/'  dataset root：$data_root/coco/a.jpg  D:\CoCo_Dataset\train2017
      num_workers = 4  # dataloader num of workers
  
      lr = 1e-4
      epoches = 20  # total epoch to train
      content_weight = 1e10  # weight of content_loss
      style_weight = 1e2  # weight of style_loss
  
      style_path = 'style_input'  # style image path
      env = 'onlyencodercontent_58_Laps_test_nores_noDynamic_10_2'  # visdom env
      plot_every = 1  # visualize in visdom for every 10 batch
  
      debug_file = '/tmp/debugnn'  # touch $debug_fie to interrupt and enter ipdb
  
      # Test Args
      content_path = 'input.png'  # input file to do style transfer [for test]
      result_path = 'output.png'  # style transfer result [for test]
  
      option_unpool = 'sum'
      cpu = False
      transfer_at_encoder = True
      transfer_at_decoder = True
      transfer_at_skip = True
      verbose = True
      save_dir = './onlyencodercontent/nores_noDynamic/58_LapSobGaus_experiments_10_2'
      log_dir = './onlyencodercontent/nores_noDynamic/58_LapSobGaus_logs_10_2'
  
      lr_decay = 5e-5
  
  
  def adjust_learning_rate(lr ,optimizer, iteration_count, lr_decay):
      """Imitating the original implementation"""
      lr = lr / (1.0 + lr_decay * iteration_count)
      for param_group in optimizer.param_groups:
          param_group['lr'] = lr
  
  def train(**kwargs):
      opt = Config()
      for k_, v_ in kwargs.items():
          setattr(opt, k_, v_)
  
      device = 'cpu' if opt.cpu or not t.cuda.is_available() else 'cuda:0'
      device = t.device(device)
      # device=t.device('cuda') if opt.use_gpu else t.device('cpu')
      vis = utils_.Visualizer(opt.env)
  
      save_dir = Path(opt.save_dir)
      save_dir.mkdir(exist_ok=True, parents=True)
      log_dir = Path(opt.log_dir)
      log_dir.mkdir(exist_ok=True, parents=True)
      writer = SummaryWriter(log_dir=str(log_dir))
      # Data loading
      transfroms = tv.transforms.Compose([
          tv.transforms.Resize(opt.image_size),
          tv.transforms.CenterCrop(opt.image_size),
          tv.transforms.ToTensor(),
          #tv.transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
          #tv.transforms.Lambda(lambda x: x*255)
      ])
      dataset = tv.datasets.ImageFolder(opt.data_root, transfroms)
      dataloader = data.DataLoader(dataset, opt.batch_size)
  
      # style transformer network
      # transformer = TransformerNet()
      print('come!')
      # visualizer = Visualizer(config)  # create a visualizer that display/save images and plots
      # device = 'cpu' if opt.cpu or not t.cuda.is_available() else 'cuda:0'
      # device = t.device(device)
  
      transfer_at = set()
      if opt.transfer_at_encoder:
          transfer_at.add('encoder')
      if opt.transfer_at_decoder:
          transfer_at.add('decoder')
      if opt.transfer_at_skip:
          transfer_at.add('skip')
      # save_dir = Path(config.save_dir)
      # save_dir.mkdir(exist_ok=True, parents=True)
      # log_dir = Path(config.log_dir)
      # log_dir.mkdir(exist_ok=True, parents=True)
      # writer = SummaryWriter(log_dir=str(log_dir))
      # vgg = net.vgg
      wct2 = Lap_Sob_Gaus(transfer_at=transfer_at, option_unpool=opt.option_unpool, device=device,
                          verbose=False)
  
      encoder = Ovodus_Laplace_model.Lap_Sob_GausEncoder(opt.option_unpool).to(device)
      decoder = Ovodus_Laplace_model.Lap_Sob_GausDecoder(opt.option_unpool).to(device)
      # vgg.load_state_dict(torch.load(config.vgg))
      # vgg = nn.Sequential(*list(vgg.children())[:31])
      laps = Lap_Sob_Gaus(transfer_at=transfer_at, option_unpool='sum', device=device)
      network = net.Net(encoder, decoder)
      network.train()
      network.to(device)
      transformer = network
      if opt.model_path:
          transformer.load_state_dict(t.load(opt.model_path, map_location=lambda _s, _: _s))
      transformer.to(device)
  
      # Vgg16 for Perceptual Loss
      # vgg = Vgg16().eval()
      # vgg.to(device)
      # for param in vgg.parameters():
      #     param.requires_grad = False
  
      # Optimizer
      # optimizer = t.optim.Adam(transformer.parameters(), opt.lr)
      enoptimizer = t.optim.Adam(network.encoder.parameters(), lr=opt.lr, betas=(0.9, 0.999))
      deoptimizer = t.optim.Adam(network.decoder.parameters(), lr=opt.lr, betas=(0.9, 0.999))
  
      # # Get style image
      # style_dataloader = utils_.get_style_data(opt.style_path, opt.batch_size)
      # #style_list = list(enumerate(style_dataloader))
      # for ii, (style, _) in tqdm.tqdm(enumerate(style_dataloader)):
      #     #a = style
      #     style = style.expand(opt.batch_size, 3, 256, 256)
      #     vis.img('style', (style.data[0] * 0.225 + 0.45).clamp(min=0, max=1))
      #     #style_list.append(style)
      #
      # style = style.to(device)
      # #
      # # #
      # # # # gram matrix for style image
      # with t.no_grad():
      #     features_style = vgg(style)
      #     gram_style = [utils_.gram_matrix(y) for y in features_style]
  
      # Loss meter
      style_meter = tnt.meter.AverageValueMeter()
      content_meter = tnt.meter.AverageValueMeter()
  
      for epoch in range(opt.epoches):
          # for jj, (style, _) in tqdm.tqdm(enumerate(style_dataloader)):
          #     a = style
          #     vis.img('style', (style.data[0] * 0.225 + 0.45).clamp(min=0, max=1))
          #     style = style.to(device)
  
          #
  
  
          content_meter.reset()
          style_meter.reset()
          for ii, (x, _) in tqdm.tqdm(enumerate(dataloader)):
              if epoch == 0:
                  adjust_learning_rate(opt.lr, enoptimizer, iteration_count=ii,  lr_decay=opt.lr_decay)
                  adjust_learning_rate(opt.lr, deoptimizer, iteration_count=ii, lr_decay=opt.lr_decay)
                  print(opt.lr)
              # style = style_list[ii][1][0]
              # # style = style_list[ii]
              # style = style.to(device)
              # # # gram matrix for style image
              # with t.no_grad():
              #     features_style = vgg(style)
              #     gram_style = [utils_.gram_matrix(y) for y in features_style]
              style_dataloader = utils_.get_style_data(opt.style_path, opt.batch_size)
              # style_list = list(enumerate(style_dataloader))
              for jj, (style, _) in tqdm.tqdm(enumerate(style_dataloader)):
                  # a = style
                  style = style.expand(opt.batch_size, 3, 256, 256)
                  #vis.img('style', (style.data[0] * 0.225 + 0.45).clamp(min=0, max=1))
                  vis.img('style', (style.data[0]).clamp(min=0, max=1))
                  # style_list.append(style)
  
              style = style.to(device)
              #
              # #
              # # # gram matrix for style image
              # with t.no_grad():
              #     features_style = vgg(style)
              #     gram_style = [utils_.gram_matrix(y) for y in features_style]
              # Train
              enoptimizer.zero_grad()
              deoptimizer.zero_grad()
              x = x.to(device)
              #y = network(x, style, Laps=laps)
              # if (ii + 1) % 10 == 0:
              # print(y)
              # y = y.clamp_(0, 1) * 255
              #y = utils_.normalize_batch(y)
              #x = utils_.normalize_batch(x)
  
              # features_y = vgg(y)
              # features_x = vgg(x)
  
              # # content loss
              # content_loss = opt.content_weight * F.mse_loss(features_y.relu2_2, features_x.relu2_2)
              #
              # # style loss
              # style_loss = 0
              #
              # for ft_y, gm_s in zip(features_y, gram_style):
              #     gram_y = utils_.gram_matrix(ft_y)
              #     style_loss += F.mse_loss(gram_y, gm_s.expand_as(gram_y))
              y, content_feats, content_loss, style_loss = network(x, style, Laps=laps)
              content_loss *= opt.content_weight
              style_loss *= opt.style_weight
              total_loss = content_loss + style_loss
              total_loss.backward()
              enoptimizer.step()
              deoptimizer.step()
  
              # Loss smooth for visualization
              content_meter.add(content_loss.item())
              style_meter.add(style_loss.item())
  
              if ii % 50 == 1:
                  print('\n')
                  print('iters:', ii, 'total_loss:', total_loss, 'loss_c:', content_loss, 'loss_s: ', style_loss)
              if (ii + 1) % opt.plot_every == 0:
                  if os.path.exists(opt.debug_file):
                      ipdb.set_trace()
  
                  # visualization
                  vis.plot('content_loss', content_meter.value()[0])
                  vis.plot('style_loss', style_meter.value()[0])
                  # denorm input/output, since we have applied (utils.normalize_batch)
                  vis.img('output1', (y.data.cpu()[0]).clamp(min=0, max=1))
                  vis.img('input1', (x.data.cpu()[0]).clamp(min=0, max=1))
                  vis.img('decoder_1', (content_feats['decoder'][0][0].data.cpu()[0]).clamp(min=0, max=1))
                  vis.img('decoder_2', (content_feats['decoder'][1][0].data.cpu()[0]).clamp(min=0, max=1))
                  vis.img('decoder_3', (content_feats['decoder'][2][0].data.cpu()[0]).clamp(min=0, max=1))
                  vis.img('decoder_4', (content_feats['decoder'][3][0].data.cpu()[0]).clamp(min=0, max=1))
                  #save_image(content_feat.clamp_(0, 1), fname_output + "decoder{:d}".format(level), padding=0)
  
              if (ii) % 1000 == 0:
                  if not os.path.exists(save_dir /'epoch_{:d}'.format(epoch)):
                      os.makedirs(save_dir /'epoch_{:d}'.format(epoch))
                  de_state_dict = network.decoder.state_dict()
                  en_state_dict = network.encoder.state_dict()
                  for key in de_state_dict.keys():
                      de_state_dict[key] = de_state_dict[key].to(t.device('cpu'))
                  t.save(de_state_dict, save_dir /'epoch_{:d}'.format(epoch)/
                         'decoder_iter_{:d}.pth.tar'.format(ii + 1))
                  for key in en_state_dict.keys():
                      en_state_dict[key] = en_state_dict[key].to(t.device('cpu'))
                  t.save(en_state_dict, save_dir /'epoch_{:d}'.format(epoch)/
                         'encoder_iter_{:d}.pth.tar'.format(ii + 1))
  
          de_state_dict = network.decoder.state_dict()
          en_state_dict = network.encoder.state_dict()
          for key in de_state_dict.keys():
              de_state_dict[key] = de_state_dict[key].to(t.device('cpu'))
          t.save(de_state_dict, save_dir /
                 'epoch_decoder_iter_{:d}.pth.tar'.format(epoch + 1))
          for key in en_state_dict.keys():
              en_state_dict[key] = en_state_dict[key].to(t.device('cpu'))
          t.save(en_state_dict, save_dir /
                 'epoch_encoder_iter_{:d}.pth.tar'.format(epoch + 1))
          # save checkpoints
          vis.save([opt.env])
          t.save(network.state_dict(), 'checkpoints/epoch_%s_style.pth' % epoch)
      writer.close()
  
  
  @t.no_grad()
  def stylize(**kwargs):
      """
      perform style transfer
      """
      opt = Config()
  
      for k_, v_ in kwargs.items():
          setattr(opt, k_, v_)
      device = t.device('cuda') if opt.use_gpu else t.device('cpu')
  
      # input image preprocess
      content_image = tv.datasets.folder.default_loader(opt.content_path)
      content_transform = tv.transforms.Compose([
          tv.transforms.ToTensor(),
          tv.transforms.Lambda(lambda x: x.mul(255))
      ])
      content_image = content_transform(content_image)
      content_image = content_image.unsqueeze(0).to(device).detach()
  
      # model setup
      style_model = TransformerNet().eval()
      style_model.load_state_dict(t.load(opt.model_path, map_location=lambda _s, _: _s))
      style_model.to(device)
  
      # style transfer and save output
      output = style_model(content_image)
      output_data = output.cpu().data[0]
      tv.utils.save_image(((output_data / 255)).clamp(min=0, max=1), opt.result_path)
  
  
  if __name__ == '__main__':
      import fire
  
      fire.Fire()
      train()
apis: torch.utils.data.DataLoader torch.load torch.no_grad torch.cuda.is_available torch.device