# VAE is disabled for this run. Including it in the model would cause style transfer # to/from just the ADE20K corpus AFAIK, as previous experiments with training # this model from scratch demonstrated that the size of the Bob Ross corpus is not # sufficient for the style transfer elements to really kick in. import sys; sys.path.append('../lib/SPADE-master/') from options.train_options import TrainOptions from models.pix2pix_model import Pix2PixModel from collections import OrderedDict import data from util.iter_counter import IterationCounter from util.visualizer import Visualizer from trainers.pix2pix_trainer import Pix2PixTrainer import os opt = TrainOptions() opt.D_steps_per_G = 1 opt.aspect_ratio = 1.0 opt.batchSize = 8 opt.beta1 = 0.0 opt.beta2 = 0.9 opt.cache_filelist_read = False opt.cache_filelist_write = False opt.checkpoints_dir = '/spell/checkpoints/' opt.contain_dontcare_label = True opt.continue_train = False opt.crop_size = 256 opt.dataroot = '/spell/adek20k' # data mount point opt.dataset_mode = "ade20k" opt.debug = False opt.display_freq = 100
python train.py --dataroot ./datasets/maps --name maps_cyclegan --model cycle_gan Train a pix2pix model: python train.py --dataroot ./datasets/facades --name facades_pix2pix --model pix2pix --direction BtoA See options/base_options.py and options/train_options.py for more training options. See training and test tips at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/tips.md See frequently asked questions at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/qa.md """ import time from options.train_options import TrainOptions from data import create_dataset from models import create_model from util.visualizer import Visualizer if __name__ == '__main__': opt = TrainOptions().parse() # get training options dataset = create_dataset( opt) # create a dataset given opt.dataset_mode and other options dataset_size = len(dataset) # get the number of images in the dataset. print('The number of training images = %d' % dataset_size) model = create_model( opt) # create a model given opt.model and other options model.setup( opt) # regular setup: load and print networks; create schedulers visualizer = Visualizer( opt) # create a visualizer that display/save images and plots total_iters = 0 # the total number of training iterations for epoch in range( opt.epoch_count, opt.niter + opt.niter_decay + 1
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). import time from collections import OrderedDict from options.train_options import TrainOptions from data.data_loader import CreateFaceConDataLoader from models.models import create_model import util.util as util from util.visualizer import Visualizer import os import numpy as np import torch from torch.autograd import Variable import tensorboardX import random opt = TrainOptions().parse() opt_test = TrainOptions().parse() opt_test.phase = 'val' opt_test.nThreads = 1 opt_test.batchSize = 1 opt_test.serial_batches = False opt_test.no_flip = True data_loader_test = CreateFaceConDataLoader(opt_test) dataset_test_ = data_loader_test.load_data() dataset_test = dataset_test_.dataset ''' for i, data in enumerate(dataset_test_): print(i) dataset_test.append(data) if (i > 10000): break
import os import numpy as np import torch from torch.autograd import Variable from collections import OrderedDict from subprocess import call import fractions def lcm(a,b): return abs(a * b)/fractions.gcd(a,b) if a and b else 0 from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.models import create_model import util.util as util from util.visualizer import Visualizer opt = TrainOptions().parse() iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') if opt.continue_train: try: start_epoch, epoch_iter = np.loadtxt(iter_path , delimiter=',', dtype=int) except: start_epoch, epoch_iter = 1, 0 print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter)) else: start_epoch, epoch_iter = 1, 0 opt.print_freq = lcm(opt.print_freq, opt.batchSize) if opt.debug: opt.display_freq = 1 opt.print_freq = 1 opt.niter = 1
import time from options.train_options import TrainOptions from data import CreateDataLoader from models import create_model from util.visualizer import Visualizer from util.laplotter import LossAccPlotter import pdb if __name__ == '__main__': opt = TrainOptions().parse() data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) # validate opt.phase = 'val' validation_loader = CreateDataLoader(opt) validation_dataset = validation_loader.load_data() validation_dataset_size = len(validation_loader) print('validate images = %d' % validation_dataset_size) model = create_model(opt) # has been initialized visualizer = Visualizer(opt) total_steps = 0 plotter = LossAccPlotter(save_to_filepath='./checkpoints/nyud_fcrn/') for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1): epoch_start_time = time.time() iter_data_time = time.time()
def main(): opt = TrainOptions().parse() data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) # read pix2pix/PAN moodel if opt.model == 'pix2pix': assert (opt.dataset_mode == 'aligned') from models.pix2pix_model import Pix2PixModel model = Pix2PixModel() model.initialize(opt) elif opt.model == 'pan': from models.pan_model import PanModel model = PanModel() model.initialize(opt) total_steps = 0 batch_size = opt.batchSize print_freq = opt.print_freq epoch_count = opt.epoch_count niter = opt.niter niter_decay = opt.niter_decay display_freq = opt.display_freq save_latest_freq = opt.save_latest_freq save_epoch_freq = opt.save_epoch_freq for epoch in range(epoch_count, niter + niter_decay + 1): epoch_start_time = time.time() epoch_iter = 0 for i, data in enumerate(dataset): # data --> (1, 3, 256, 256) iter_start_time = time.time() total_steps += batch_size epoch_iter += batch_size model.set_input(data) model.optimize_parameters() if total_steps % print_freq == 0: errors = model.get_current_errors() t = (time.time() - iter_start_time) / batch_size message = '(epoch: %d, iters: %d, time: %.3f) ' % ( epoch, epoch_iter, t) for k, v in errors.items(): message += '%s: %.3f ' % (k, v) print(message) # save latest weights if total_steps % save_latest_freq == 0: print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps)) model.save('latest') # save weights periodicaly if epoch % save_epoch_freq == 0: print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps)) model.save('latest') model.save(epoch) print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, niter + niter_decay, time.time() - epoch_start_time)) model.update_learning_rate()
import time from collections import OrderedDict from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.mapping_model import Pix2PixHDModel_Mapping import util.util as util from util.visualizer import Visualizer import os import numpy as np import torch import torchvision.utils as vutils from torch.autograd import Variable import datetime import random opt = TrainOptions().parse() visualizer = Visualizer(opt) iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') if opt.continue_train: try: start_epoch, epoch_iter = np.loadtxt(iter_path, delimiter=',', dtype=int) except: start_epoch, epoch_iter = 1, 0 visualizer.print_save('Resuming from epoch %d at iteration %d' % (start_epoch - 1, epoch_iter)) else: start_epoch, epoch_iter = 1, 0 if opt.which_epoch != "latest":
from models import create_model import torch import torchvision import torchvision.transforms as transforms from util import util import numpy as np import progressbar as pb import shutil import datetime as dt import matplotlib.pyplot as plt if __name__ == '__main__': opt = TrainOptions().parse() opt.load_model = True opt.num_threads = 1 # test code only supports num_threads = 1 opt.batch_size = 1 # test code only supports batch_size = 1 opt.display_id = -1 # no visdom display opt.phase = 'test' opt.dataroot = './dataset/ilsvrc2012/%s/' % opt.phase opt.loadSize = 256 opt.how_many = 5 opt.aspect_ratio = 1.0 opt.sample_Ps = [ 6, ] opt.load_model = True # number of random points to assign
import time from options.train_options import TrainOptions from dataset import image_data from models import create_model import shutil import os import torch from torch.utils.data import random_split, DataLoader import numpy as np from tqdm import tqdm if __name__ == '__main__': trainoptions = TrainOptions() opt = trainoptions.parse() # setup os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.gpu_ids) projroot = os.path.join(os.getcwd(), opt.study) modelroot = os.path.join(projroot, 'save_models') dataroot = os.path.join(projroot, 'data') modelsubfolder = '{}_{}'.format(opt.study, opt.model) modelfolder = os.path.join(modelroot, modelsubfolder) opt.name = opt.name + '_' + modelsubfolder torch.manual_seed(opt.seed) torch.cuda.manual_seed(opt.seed) opt.input_nc, opt.output_nc = 1, 1
# By Samaneh Azadi ################################################################################ from torch import transpose import torch.utils.data as data from torch import index_select, LongTensor from PIL import Image import os import os.path import numpy as np from scipy import misc import random from options.train_options import TrainOptions import torch opt = TrainOptions().parse() # set CUDA_VISIBLE_DEVICES before import torch IMG_EXTENSIONS = ['.png'] def is_image_file(filename): return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) def make_dataset(dir): images = [] assert os.path.isdir(dir), '%s is not a valid directory' % dir for root, _, fnames in sorted(os.walk(dir)): for fname in fnames: if is_image_file(fname):
visuals = OrderedDict([('synthesized_image', trainer.get_latest_generated()), ('real_image', data_i['image'])]) visualizer.display_current_results(visuals, epoch, iter_counter.total_steps_so_far) if rank == 0: print('saving the latest model (epoch %d, total_steps %d)' % (epoch, iter_counter.total_steps_so_far)) trainer.save('latest') iter_counter.record_current_iter() trainer.update_learning_rate(epoch) iter_counter.record_epoch_end() if (epoch % opt.save_epoch_freq == 0 or epoch == iter_counter.total_epochs) and (rank == 0): print('saving the model at the end of epoch %d, iters %d' % (epoch, iter_counter.total_steps_so_far)) trainer.save(epoch) print('Training was successfully finished.') if __name__ == '__main__': global TrainOptions TrainOptions = TrainOptions() opt = TrainOptions.parse(save=True) opt.world_size = opt.num_gpu opt.mpdist = True mp.set_start_method('spawn', force=True) mp.spawn(main_worker, nprocs=opt.world_size, args=(opt.world_size, opt))
python train.py --name_prefix demo --dataname RIMEScharH32W16 --capitalize --display_port 8192 See options/base_options.py and options/train_options.py for more training options. """ import time from options.train_options import TrainOptions from data import create_dataset from models import create_model from util.visualizer import Visualizer from util.util import seed_rng from util.util import prepare_z_y, get_curr_data import torch import os if __name__ == '__main__': opt = TrainOptions().parse() # get training options # Seed RNG seed_rng(opt.seed) torch.backends.cudnn.benchmark = True dataset = create_dataset( opt) # create a dataset given opt.dataset_mode and other options dataset_size = len(dataset) # get the number of images in the dataset. print('The number of training images = %d' % dataset_size) model = create_model( opt) # create a model given opt.model and other options model.setup( opt) # regular setup: load and print networks; create schedulers if opt.single_writer: opt.G_init = 'N02' opt.D_init = 'N02'
print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps)) model.save('latest') model.save(epoch) print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time)) if epoch > opt.niter: model.update_learning_rate() if __name__ == '__main__': freeze_support() # python train.py --dataroot /.path_to_your_data --learn_residual --resize_or_crop crop --fineSize CROP_SIZE (we used 256) opt = TrainOptions().parse() opt.dataroot = 'D:\Photos\TrainingData\BlurredSharp\combined' opt.learn_residual = True opt.resize_or_crop = "crop" opt.fineSize = 256 opt.gan_type = "gan" # opt.which_model_netG = "unet_256" # default = 5000 opt.save_latest_freq = 100 # default = 100 opt.print_freq = 20 data_loader = CreateDataLoader(opt) model = create_model(opt)
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved. ### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.models import create_model import os import util.util as util from torch.autograd import Variable import torch.nn as nn opt = TrainOptions().parse() opt.nThreads = 1 opt.batchSize = 1 opt.serial_batches = True opt.no_flip = True opt.instance_feat = True name = 'features' save_path = os.path.join(opt.checkpoints_dir, opt.name) ############ Initialize ######### data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) model = create_model(opt) util.mkdirs(os.path.join(opt.dataroot, opt.phase + '_feat')) ######## Save precomputed feature maps for 1024p training ####### for i, data in enumerate(dataset): print('%d / %d images' % (i+1, dataset_size)) feat_map = model.module.netE.forward(Variable(data['image'].cuda(), volatile=True), data['inst'].cuda())
import os import sys from collections import OrderedDict from options.train_options import TrainOptions import data from data.base_dataset import repair_data from util.iter_counter import IterationCounter from util.visualizer import Visualizer from trainers.pix2pix_trainer import Pix2PixTrainer from options.test_options import TestOptions import tqdm from util import html from util.util import tensor2im, tensor2label # parse options opt = TrainOptions().parse() # print options to help debugging print(' '.join(sys.argv)) # load the dataset dataloader = data.create_dataloader(opt) # create trainer for our model trainer = Pix2PixTrainer(opt) # create tool for counting iterations iter_counter = IterationCounter(opt, len(dataloader)) # create tool for visualization visualizer = Visualizer(opt)
def main(): opt = TrainOptions().parse() train_history = TrainHistory() checkpoint = Checkpoint(opt) os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_id """Architecture""" net = MakeLinearModel(1024,16) net = torch.nn.DataParallel(net).cuda() """Uploading Mean and SD""" path_to_data = '.../multi-view-pose-estimation/dataset/' #mean and sd of 2d poses in training dataset Mean_2D = np.loadtxt(path_to_data + 'Mean_2D.txt') Mean_2D = Mean_2D.astype('float32') Mean_2D = torch.from_numpy (Mean_2D) Mean_Delta = np.loadtxt(path_to_data + 'Mean_Delta.txt') Mean_Delta = Mean_Delta.astype('float32') Mean_Delta = torch.from_numpy (Mean_Delta) SD_2D = np.loadtxt(path_to_data + 'SD_2D.txt') SD_2D = SD_2D.astype('float32') SD_2D = torch.from_numpy (SD_2D) SD_Delta = np.loadtxt(path_to_data + 'SD_Delta.txt') SD_Delta = SD_Delta.astype('float32') SD_Delta = torch.from_numpy (SD_Delta) """Loading Data""" train_list = 'train_list.txt' train_loader = torch.utils.data.DataLoader( data.PtsList(train_list, is_train=True), batch_size=opt.bs, shuffle=True, num_workers=opt.nThreads, pin_memory=True) val_list = 'valid_list.txt' val_loader = torch.utils.data.DataLoader( data.PtsList(val_list,is_train=False ), batch_size=opt.bs, shuffle=False, num_workers=opt.nThreads, pin_memory=True) """optimizer""" optimizer = torch.optim.Adam( net.parameters(), lr=opt.lr, betas=(0.9,0.999), weight_decay=0) """training and validation""" for epoch in range(0, opt.nEpochs): adjust_learning_rate(optimizer, epoch, opt.lr) # train for one epoch train_loss = train(train_loader, net, Mean_2D, Mean_Delta, SD_2D, SD_Delta, optimizer, epoch, opt) # evaluate on validation set val_loss,val_pckh = validate(val_loader, net, Mean_2D, Mean_Delta, SD_2D, SD_Delta, epoch, opt) # update training history e = OrderedDict( [('epoch', epoch)] ) lr = OrderedDict( [('lr', opt.lr)] ) loss = OrderedDict( [('train_loss', train_loss),('val_loss', val_loss)] ) pckh = OrderedDict( [('val_pckh', val_pckh)] ) train_history.update(e, lr, loss, pckh) checkpoint.save_checkpoint(net, train_history, 'best-single.pth.tar')
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved. ### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.models import create_model import numpy as np import os opt = TrainOptions().parse() opt.nThreads = 1 opt.batchSize = 1 opt.serial_batches = True opt.no_flip = True opt.instance_feat = True name = 'features' save_path = os.path.join(opt.checkpoints_dir, opt.name) ############ Initialize ######### data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) model = create_model(opt) ########### Encode features ########### reencode = True if reencode: features = {} for label in range(opt.label_nc): features[label] = np.zeros((0, opt.feat_num + 1)) for i, data in enumerate(dataset):
""" import time from datasets.get_cityscapes import get_loaders_cityscapes from options.train_options import TrainOptions from data import create_dataset from models import create_model from util.evaluation import evaluate from util.transforms import joint_transform_train from util.util import calculate_class_weights from util.visualizer import Visualizer from pathlib import Path import torch if __name__ == '__main__': opt = TrainOptions().parse() # get training options # dataset = create_dataset(opt) # create a dataset given opt.dataset_mode and other options if opt.run_colab: root_path = Path("/content/data/") else: root_path = Path( "/Users/laurenssamson/Documents/Projects/data/Cityscapes/leftImg8bit_trainvaltest" ) training_loader, val_loader, train_eval_loader = get_loaders_cityscapes( root_path, opt) opt.class_weights = calculate_class_weights(val_loader, opt.output_nc) model = create_model( opt) # create a model given opt.model and other options model.setup( opt) # regular setup: load and print networks; create schedulers visualizer = Visualizer(
def train_main(raw_args=None): # print(torch.backends.cudnn.benchmark) opt = TrainOptions().parse(raw_args) # get training options if opt.debug_mode: import multiprocessing multiprocessing.set_start_method('spawn', True) opt.num_threads = 0 dataset = create_dataset( opt) # create a dataset given opt.dataset_mode and other options dataset_size = len(dataset) # get the number of images in the dataset. print('The number of training images = %d' % dataset_size) existing_epochs = glob.glob(opt.checkpoints_dir + '/' + opt.name + '/*[0-9]_net_G_A.pth') if opt.restart_training and len(existing_epochs) > 0: opt.epoch = int( os.path.splitext(os.path.basename( existing_epochs[-1]))[0].split('_')[0]) opt.epoch_count = opt.epoch + 1 plot_losses_from_log_files(opt, dataset_size, domain=['A', 'B'], specified=['G', 'D', 'cycle']) model = create_model( opt) # create a model given opt.model and other options model.setup( opt) # regular setup: load and print networks; create schedulers visualizer = Visualizer( opt) # create a visualizer that display/save images and plots total_iters = 0 # the total number of training iterations for epoch in range( opt.epoch_count, opt.niter + opt.niter_decay + 1 ): # outer loop for different epochs; we save the model by <epoch_count>, <epoch_count>+<save_latest_freq> epoch_start_time = time.time() # timer for entire epoch iter_data_time = time.time() # timer for data loading per iteration epoch_iter = 0 # the number of training iterations in current epoch, reset to 0 every epoch for i, data in enumerate(dataset): # inner loop within one epoch iter_start_time = time.time( ) # timer for computation per iteration if total_iters % opt.print_freq == 0: t_data = iter_start_time - iter_data_time visualizer.reset() total_iters += opt.batch_size epoch_iter += opt.batch_size model.set_input( data) # unpack data from dataset and apply preprocessing model.optimize_parameters( ) # calculate loss functions, get gradients, update network weights if total_iters % opt.display_freq == 0: # display images on visdom and save images to a HTML file save_result = total_iters % opt.update_html_freq == 0 model.compute_visuals() visualizer.display_current_results(model.get_current_visuals(), epoch, save_result) if total_iters % opt.print_freq == 0: # print training losses and save logging information to the disk losses = model.get_current_losses() t_comp = (time.time() - iter_start_time) / opt.batch_size visualizer.print_current_losses(epoch, epoch_iter, losses, t_comp, t_data) if opt.display_id > 0: visualizer.plot_current_losses( epoch, float(epoch_iter) / dataset_size, losses) if total_iters % opt.save_latest_freq == 0: # cache our latest model every <save_latest_freq> iterations print('saving the latest model (epoch %d, total_iters %d)' % (epoch, total_iters)) save_suffix = 'iter_%d' % total_iters if opt.save_by_iter else 'latest' model.save_networks(save_suffix) iter_data_time = time.time() if epoch % opt.save_epoch_freq == 0: # cache our model every <save_epoch_freq> epochs print('saving the model at the end of epoch %d, iters %d' % (epoch, total_iters)) model.save_networks('latest') model.save_networks(epoch) print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time)) model.update_learning_rate( ) # update learning rates at the end of every epoch.
# Generates a ton of images and ranks them in order of # discriminator loss def plotTensor(im): im = util.tensor2im(im) plotim(im) def plotim(im): plt.imshow(im) plt.show() if __name__ == '__main__': opt = TrainOptions().parse() opt.no_flip = True opt.resize_or_crop = 'none' opt.dataset_mode = 'auto' data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) dirname = '12_nol1' model = create_model(opt) model.setup(opt) total_steps = 0 chkpt_D = torch.load('checkpoints/sv_nlayers5_ranker/2_net_D.pth')
def main(): plt.ion() # 开启interactive mode,便于连续plot opt = TrainOptions().parse() # 用于计算的设备 CPU or GPU device = torch.device("cuda" if USE_CUDA else "cpu") # 定义判别器与生成器的网络 #net_d = NLayerDiscriminator(opt.output_nc, opt.ndf, n_layers=3)#batchnorm #net_d = Discriminator(opt.output_nc) net_d_ct =networks.define_D(opt.output_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, opt.gpu_ids) net_d_dr =networks.define_D(opt.input_nc, opt.ndf, 'ProjNet', opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, opt.gpu_ids) net_g_dr=networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain, opt.gpu_ids) #net_g = CTGenerator(opt.input_nc, opt.output_nc, opt.ngf, n_blocks=6) net_g_ct = networks.define_G(1, 65, opt.ngf, 'CTnet', opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain, opt.gpu_ids) # init_weights(net_d_dr) # init_weights(net_d_ct) # init_weights(net_g_dr) # init_weights(net_g_ct) net_d_ct.to(device) net_d_dr.to(device) net_g_dr.to(device) net_g_ct.to(device) one = torch.FloatTensor([1]) mone = one * -1 one = one.to(device) mone= mone.to(device) #summary(net_g_dr, (2,65, 65,65)) if load_net: # save_filename = 'net_d%s.pth' % epoch_start # save_path = os.path.join('./check/', save_filename) # load_network(net_d, save_path) save_filename = 'net_g%s.pth' % epoch_start save_path = os.path.join('./check/', save_filename) load_network(net_g_ct, save_path) # 损失函数 #criterion = nn.BCELoss().to(device) criterion = nn.MSELoss().to(device) criterion1 = nn.L1Loss().to(device) # 优化器 optimizer_d = torch.optim.Adam(itertools.chain(net_d_ct.parameters(),net_d_dr.parameters()), lr=0.0001,betas=[0.5,0.9]) optimizer_g = torch.optim.Adam(itertools.chain(net_g_ct.parameters(),net_g_dr.parameters()), lr=0.0001,betas=[0.5,0.9]) #optimizer_d = torch.optim.AdamW(net_d.parameters(), lr=0.0001) #optimizer_g = torch.optim.AdamW(net_g.parameters(), lr=0.0001) #one = torch.FloatTensor([1]).cuda() #mone = one * -1 dataset = create_dataset(opt) # create a dataset given opt.dataset_mode and other options dataset_size = len(dataset) # get the number of images in the dataset. def gensample(): for image in enumerate(dataset): yield image gen=gensample() ii=0 for epoch in range(MAX_EPOCH): # 为真实数据加上噪声 for it,data in enumerate(dataset): #载入数据 dr_real = autograd.Variable(data['A'].cuda()) dr_real=dr_real.squeeze(0) ct_real = autograd.Variable(data['B'].cuda()) ct_real=ct_real.squeeze(0) # 训练 #内循环 freeze_params(net_g_ct) freeze_params(net_g_dr) unfreeze_params(net_d_ct) unfreeze_params(net_d_dr) ct_fake = autograd.Variable(net_g_ct(dr_real).data) dr_fake = autograd.Variable(net_g_dr(ct_real).data) optimizer_d.zero_grad() loss_dsc_realct = net_d_ct(ct_real).mean() #loss_dsc_realct.backward() loss_dsc_fakect = net_d_ct(ct_fake.detach()).mean() #loss_dsc_fakect.backward() gradient_penalty_ct = calc_gradient_penalty(net_d_ct, ct_real, ct_fake) #gradient_penalty_ct.backward() loss_d_ct=loss_dsc_fakect - loss_dsc_realct+gradient_penalty_ct loss_d_ct.backward() Wd_ct=loss_dsc_realct-loss_dsc_fakect loss_dsc_realdr = net_d_dr(dr_real).mean() #loss_dsc_realdr.backward() loss_dsc_fakedr = net_d_dr(dr_fake.detach()).mean() #loss_dsc_fakedr.backward() gradient_penalty_dr = calc_gradient_penalty(net_d_dr, dr_real, dr_fake) #gradient_penalty_dr.backward() loss_d_dr = loss_dsc_fakedr - loss_dsc_realdr + gradient_penalty_dr loss_d_dr.backward() Wd_dr = loss_dsc_realdr - loss_dsc_fakedr optimizer_d.step() if it%CRITIC_ITERS==0: #if True: unfreeze_params(net_g_ct) freeze_params(net_d_ct) unfreeze_params(net_g_dr) freeze_params(net_d_dr) ct_fake_g=net_g_ct(dr_real) dr_fake_g=net_g_dr(ct_real) #外循环ct_dr # optimizer_g.zero_grad() loss_out_dr=criterion1(net_g_ct(dr_fake_g),ct_real) # loss_out_dr.backward() # optimizer_g.step() #net_g_ct.load_state_dict(dict_g_ct) #外循环dr_ct # optimizer_g.zero_grad() loss_out_ct=criterion1(net_g_dr(ct_fake_g),dr_real) # loss_out_ct.backward() # optimizer_g.step() #内循环gan loss_g_ct = - net_d_ct(ct_fake_g).mean() #loss_g_ct.backward() loss_g_dr = - net_d_dr(dr_fake_g).mean() #loss_g_dr.backward() loss_gan = loss_out_dr + loss_out_ct #loss_gan=loss_out_dr+loss_out_ct+loss_g_ct+loss_g_dr #loss_gan = criterion(net_g_ct(dr_real), ct_real) + criterion(net_g_dr(ct_real), dr_real) optimizer_g.zero_grad() loss_gan.backward() optimizer_g.step() if it%1==0: fk_im=toimage(torch.irfft(torch.roll(torch.roll(ct_fake,-32,2),-32,3).permute(1, 2, 3, 0), 2, onesided=False)[32, :, :].unsqueeze(0)) #fk_im=toimage(ct_fake[0,32,:,:].unsqueeze(0)) #img_test.append(fk_im) save_filenamet = 'fakect%s.bmp' % int(epoch/dataset_size) img_path = os.path.join('./check/img/', save_filenamet) save_image(fk_im, img_path) rel_im=toimage(torch.irfft(torch.roll(torch.roll(ct_real,-32,2),-32,3).permute(1, 2, 3, 0), 2, onesided=False)[32, :, :].unsqueeze(0)) #rel_im = toimage(ct_real[0,32, :, :].unsqueeze(0)) # img_test.append(rel_im) save_image(rel_im, os.path.join('./check/img/', 'Realct%s.bmp' % int(epoch))) fake_im = toimage(torch.irfft(torch.roll(torch.roll(dr_fake, -128, 2), -128, 3).permute(1, 2, 3, 0), 2, onesided=False)) #fake_im =toimage(dr_fake.squeeze(0)) save_image(fake_im, os.path.join('./check/img/', 'fakedr%s.bmp' % int(epoch))) ceshi(net_g_ct) message = '(epoch: %d, iters: %d, D_ct: %.3f;[real:%.3f;fake:%.3f], G_ct: %.3f, D_dr: %.3f, G_dr: %.3f) ' % (int(epoch), ii,loss_d_ct,loss_dsc_realct,loss_dsc_fakect,loss_g_ct,loss_d_dr,loss_g_dr) print(message) save_filename = 'net_g%s.pth' % epoch save_path = os.path.join('./check/', save_filename) torch.save(net_g_ct.cpu().state_dict(), save_path) net_g_ct.cuda(0)
from options.train_options import TrainOptions from data import create_dataset from models import create_model from util.visualizer import Visualizer from util import util from math import log10 from util.visualizer import save_sr_result import time import os import pytorch_ssim if __name__ == '__main__': # 加载设置 opt = TrainOptions().parse() opt.upscale_factor = 4 # 设置显示验证结果存储的设置 web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'val') image_dir = os.path.join(web_dir, 'images') util.mkdirs([web_dir, image_dir]) # 加载训练数据集 dataset_train = create_dataset(opt) dataset_train_size = len(dataset_train) print('The number of training images = %d' % dataset_train_size) # 加载验证数据集 opt1 = TrainOptions().parse() opt1.upscale_factor = 4 opt1.phase = "val" opt1.batch_size = 1
from data.data_loader import CreateDataLoader import util.util as util from util.visualizer import Visualizer import os import numpy as np import torch import torchvision import torchvision.transforms as transforms import scipy.io as sio # Set random seed torch.manual_seed(0) np.random.seed(0) # Extract the options opt = TrainOptions().parse() # For testing the neural networks, manually edit/add options below opt.gan_mode = 'none' # 'wgangp', 'lsgan', 'vanilla', 'none' # Set the input dataset opt.dataset_mode = 'CIFAR10' # Current dataset: CIFAR10, CelebA if opt.dataset_mode in ['CIFAR10', 'CIFAR100']: opt.n_layers_D = 3 opt.label_smooth = 1 # Label smoothing factor (for lsgan and vanilla gan only) opt.n_downsample = 2 # Downsample times opt.n_blocks = 2 # Numebr of residual blocks opt.first_kernel = 5 # The filter size of the first convolutional layer in encoder opt.batchsize = 128 opt.n_epochs = 200 # # of epochs without lr decay
def main(): cfg = TrainOptions().parse() # get training options cfg.NUM_GPUS = torch.cuda.device_count() cfg.batch_size = int(cfg.batch_size / max(1, cfg.NUM_GPUS)) cfg.phase = 'train' launch_job(cfg=cfg, init_method=cfg.init_method, func=train)
def __len__(self): return len(self.path_to_images) def __getitem__(self, idx): idx = idx.tolist() if torch.is_tensor(idx) else idx img_name = self.path_to_images[idx] image = PIL.open(img_name) if self.transform: sample = self.transform(sample) return sample if __name__ == '__main__': opt = TrainOptions().parse() opt.dataroot = './dataset/ilsvrc2012/%s/' % opt.phase colorization_dataset = Colorization_Dataset(opt.dataroot, transform=transforms.Compose([ transforms.RandomChoice([transforms.Resize(opt.loadSize, interpolation=1), transforms.Resize(opt.loadSize, interpolation=2), transforms.Resize(opt.loadSize, interpolation=3), transforms.Resize((opt.loadSize, opt.loadSize), interpolation=1), transforms.Resize((opt.loadSize, opt.loadSize), interpolation=2), transforms.Resize((opt.loadSize, opt.loadSize), interpolation=3)]), transforms.RandomChoice([transforms.RandomResizedCrop(opt.fineSize, interpolation=1), transforms.RandomResizedCrop(opt.fineSize, interpolation=2), transforms.RandomResizedCrop(opt.fineSize, interpolation=3)]), transforms.RandomHorizontalFlip(), transforms.ToTensor()]))
# ======================================================== # Compositional GAN # Train different components of the paired/unpaired models # By Samaneh Azadi # ======================================================== import time from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.models import create_model from util.visualizer import Visualizer opt = TrainOptions().parse() opt.phase = 'train' data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) model = create_model(opt) visualizer = Visualizer(opt) print('Train the STN models') # Only for the unpaired case if opt.dataset_mode == 'comp_decomp_unaligned' and opt.niterSTN: opt.isPretrain = False visualizer = Visualizer(opt) total_steps = 0 for epoch in range(opt.epoch_count, opt.niterSTN + 1): epoch_start_time = time.time() epoch_iter = 0
See options/base_options.py and options/train_options.py for more training options. See training and test tips at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/tips.md See frequently asked questions at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/qa.md """ import time from options.train_options import TrainOptions from options.val_options import ValOptions from data import create_dataset from models import create_model from util.visualizer_sssd import Visualizer, save_images, cal_scores from copy import deepcopy import os from util import html if __name__ == '__main__': opt = TrainOptions().parse() # get training options # opt.sample_nums = 100 # opt.dataset_mode = "alignedm2md" dataset = create_dataset( opt) # create a dataset given opt.dataset_mode and other options dataset_size = len(dataset) # get the number of images in the dataset. print('The number of training images = %d' % dataset_size) model = create_model( opt) # create a model given opt.model and other options model.setup( opt) # regular setup: load and print networks; create schedulers visualizer = Visualizer( opt) # create a visualizer that display/save images and plots total_iters = 0 # the total number of training iterations
import time from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.models import create_model opt = TrainOptions().parse() data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) model = create_model(opt) total_steps = 0 for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1): epoch_start_time = time.time() epoch_iter = 0 for i, data in enumerate(dataset): iter_start_time = time.time() total_steps += opt.batchSize epoch_iter += opt.batchSize model.set_input(data) model.optimize_parameters() if total_steps % opt.print_freq == 0: errors = model.get_current_errors() t = (time.time() - iter_start_time) / opt.batchSize
def main(): opt = TrainOptions() args = opt.initialize() _t = {'iter time': Timer()} model_name = args.source + '_to_' + args.target if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) os.makedirs(os.path.join(args.snapshot_dir, 'logs')) opt.print_options(args) sourceloader, targetloader = CreateSrcDataLoader( args), CreateTrgDataLoader(args) targetloader_iter, sourceloader_iter = iter(targetloader), iter( sourceloader) model, optimizer = CreateModel(args) model_D, optimizer_D = CreateDiscriminator(args) start_iter = 0 if args.restore_from is not None: start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1]) train_writer = tensorboardX.SummaryWriter( os.path.join(args.snapshot_dir, "logs", model_name)) bce_loss = torch.nn.BCEWithLogitsLoss() cudnn.enabled = True cudnn.benchmark = True model.train() model.cuda() model_D.train() model_D.cuda() loss = [ 'loss_seg_src', 'loss_seg_trg', 'loss_D_trg_fake', 'loss_D_src_real', 'loss_D_trg_real' ] _t['iter time'].tic() for i in range(start_iter, args.num_steps): model.adjust_learning_rate(args, optimizer, i) model_D.adjust_learning_rate(args, optimizer_D, i) optimizer.zero_grad() optimizer_D.zero_grad() for param in model_D.parameters(): param.requires_grad = False src_img, src_lbl, _, _ = sourceloader_iter.next() src_img, src_lbl = Variable(src_img).cuda(), Variable( src_lbl.long()).cuda() src_seg_score = model(src_img, lbl=src_lbl) loss_seg_src = model.loss loss_seg_src.backward() if args.data_label_folder_target is not None: trg_img, trg_lbl, _, _ = targetloader_iter.next() trg_img, trg_lbl = Variable(trg_img).cuda(), Variable( trg_lbl.long()).cuda() trg_seg_score = model(trg_img, lbl=trg_lbl) loss_seg_trg = model.loss else: trg_img, _, name = targetloader_iter.next() trg_img = Variable(trg_img).cuda() trg_seg_score = model(trg_img) loss_seg_trg = 0 outD_trg = model_D(F.softmax(trg_seg_score), 0) loss_D_trg_fake = model_D.loss loss_trg = args.lambda_adv_target * loss_D_trg_fake + loss_seg_trg loss_trg.backward() for param in model_D.parameters(): param.requires_grad = True src_seg_score, trg_seg_score = src_seg_score.detach( ), trg_seg_score.detach() outD_src = model_D(F.softmax(src_seg_score), 0) loss_D_src_real = model_D.loss / 2 loss_D_src_real.backward() outD_trg = model_D(F.softmax(trg_seg_score), 1) loss_D_trg_real = model_D.loss / 2 loss_D_trg_real.backward() optimizer.step() optimizer_D.step() for m in loss: train_writer.add_scalar(m, eval(m), i + 1) if (i + 1) % args.save_pred_every == 0: print('taking snapshot ...') torch.save( model.state_dict(), os.path.join(args.snapshot_dir, '%s_' % (args.source) + str(i + 1) + '.pth')) if (i + 1) % args.print_freq == 0: _t['iter time'].toc(average=False) print('[it %d][src seg loss %.4f][lr %.4f][%.2fs]' % \ (i + 1, loss_seg_src.data, optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff)) if i + 1 > args.num_steps_stop: print('finish training') break _t['iter time'].tic()
import time from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.models import create_model import ntpath import os from util import util import shutil import numpy as np #from util.visualizer import Visualizer opt = TrainOptions().parse() # Load dataset_train data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) # Load dataset_test opt.batchSize = 1 opt.phase = 'test' data_loader_test = CreateDataLoader(opt) dataset_test = data_loader_test.load_data() dataset_size = len(data_loader_test) print('#testing images = %d' % dataset_size) # Create or clear dir for saving generated samples if os.path.exists(opt.testing_path): shutil.rmtree(opt.testing_path)
def main(): opt = TrainOptions().parse() # Determine validation step options that might differ from training if opt.data == 'KTH': val_pick_mode = 'Slide' val_gpu_ids = [opt.gpu_ids[0]] val_batch_size = 1 elif opt.data in ['UCF', 'HMDB51', 'S1M']: val_pick_mode = 'First' val_gpu_ids = opt.gpu_ids val_batch_size = opt.batch_size / 2 else: raise ValueError('Dataset [%s] not recognized.' % opt.data) expr_dir = os.path.join(opt.checkpoints_dir, opt.name) makedir(expr_dir) tb_dir = os.path.join(opt.tensorboard_dir, opt.name) makedir(tb_dir) file_name = os.path.join(expr_dir, 'train_opt.txt') with open(file_name, 'wt') as opt_file: listopt(opt, opt_file) log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt') print('after reading options') include_following = (opt.model_type != 'mcnet') data_loader = CustomDataLoader(opt.data, opt.c_dim, opt.dataroot, opt.textroot, opt.video_list, opt.K, opt.T, opt.backwards, opt.flip, opt.pick_mode, opt.image_size, include_following, opt.skip, opt.F, opt.batch_size, opt.serial_batches, opt.nThreads) print(data_loader.name()) dataset = data_loader.load_data() dataset_size = len(data_loader) print('# training videos = %d' % dataset_size) env = create_environment( opt.model_type, opt.gf_dim, opt.c_dim, opt.gpu_ids, True, opt.checkpoints_dir, opt.name, opt.K, opt.T, opt.F, opt.image_size, opt.batch_size, opt.which_update, opt.comb_type, opt.shallow, opt.ks, opt.num_block, opt.layers, opt.kf_dim, opt.enable_res, opt.rc_loc, opt.no_adversarial, opt.alpha, opt.beta, opt.D_G_switch, opt.margin, opt.lr, opt.beta1, opt.sn, opt.df_dim, opt.Ip, opt.continue_train, opt.comb_loss) total_updates = env.start_update writer = SummaryWriter(log_dir=os.path.join(opt.tensorboard_dir, opt.name)) while True: for data in dataset: iter_start_time = time.time() # Enable losses on intermediate and final predictions partway through training if total_updates >= opt.inter_sup_update: env.enable_inter_loss() if total_updates >= opt.final_sup_update: env.enable_final_loss() # Update model total_updates += 1 env.set_inputs(data) env.optimize_parameters() if total_updates % opt.print_freq == 0: errors = env.get_current_errors() t = (time.time() - iter_start_time) / opt.batch_size writer.add_scalar('iter_time', t, total_updates) for key in errors.keys(): writer.add_scalar('loss/%s' % (key), errors[key], total_updates) print_current_errors(log_name, total_updates, errors, t) if total_updates % opt.display_freq == 0: visuals = env.get_current_visuals() grid = visual_grid(visuals, opt.K, opt.T) writer.add_image('current_batch', grid, total_updates) if total_updates % opt.save_latest_freq == 0: print('saving the latest model (update %d)' % total_updates) env.save('latest', total_updates) env.save(total_updates, total_updates) if total_updates % opt.validate_freq == 0: psnr_plot, ssim_plot, grid = val( opt.c_dim, opt.data, opt.T * 2, opt.dataroot, opt.textroot, 'val_data_list.txt', opt.K, opt.backwards, opt.flip, val_pick_mode, opt.image_size, val_gpu_ids, opt.model_type, opt.skip, opt.F, val_batch_size, True, opt.nThreads, opt.gf_dim, False, opt.checkpoints_dir, opt.name, opt.no_adversarial, opt.alpha, opt.beta, opt.D_G_switch, opt.margin, opt.lr, opt.beta1, opt.sn, opt.df_dim, opt.Ip, opt.comb_type, opt.comb_loss, opt.shallow, opt.ks, opt.num_block, opt.layers, opt.kf_dim, opt.enable_res, opt.rc_loc, opt.continue_train, 'latest') writer.add_image('psnr', psnr_plot, total_updates) writer.add_image('ssim', ssim_plot, total_updates) writer.add_image('samples', grid, total_updates) if total_updates >= opt.max_iter: env.save('latest', total_updates) break if total_updates >= opt.max_iter: break
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved. ### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). import time from collections import OrderedDict from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.models import create_model import util.util as util from util.visualizer import Visualizer import os import numpy as np import torch from torch.autograd import Variable opt = TrainOptions().parse() iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') if opt.continue_train: try: start_epoch, epoch_iter = np.loadtxt(iter_path , delimiter=',', dtype=int) except: start_epoch, epoch_iter = 1, 0 print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter)) else: start_epoch, epoch_iter = 1, 0 if opt.debug: opt.display_freq = 1 opt.print_freq = 1 opt.niter = 1 opt.niter_decay = 0 opt.max_dataset_size = 10
def main(): opt = TrainOptions().parse() if opt.sr_dir == '': print('sr directory is null.') exit() sr_pretrain_dir = os.path.join( opt.exp_dir, opt.exp_id, opt.sr_dir + '-' + opt.load_prefix_pose[0:-1]) if not os.path.isdir(sr_pretrain_dir): os.makedirs(sr_pretrain_dir) train_history = ASNTrainHistory() # print(train_history.lr) # exit() checkpoint_agent = Checkpoint() visualizer = Visualizer(opt) visualizer.log_path = sr_pretrain_dir + '/' + 'log.txt' train_scale_path = sr_pretrain_dir + '/' + 'train_scales.txt' train_rotation_path = sr_pretrain_dir + '/' + 'train_rotations.txt' val_scale_path = sr_pretrain_dir + '/' + 'val_scales.txt' val_rotation_path = sr_pretrain_dir + '/' + 'val_rotations.txt' # with open(visualizer.log_path, 'a+') as log_file: # log_file.write(opt.resume_prefix_pose + '.pth.tar\n') # lost_joint_count_path = os.path.join(opt.exp_dir, opt.exp_id, opt.astn_dir, 'joint-count.txt') # print("=> log saved to path '{}'".format(visualizer.log_path)) # if opt.dataset == 'mpii': # num_classes = 16 os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_id print('collecting training scale and rotation distributions ...\n') train_scale_distri = read_grnd_distri_from_txt(train_scale_path) train_rotation_distri = read_grnd_distri_from_txt(train_rotation_path) dataset = MPII('dataset/mpii-hr-lsp-normalizer.json', '/bigdata1/zt53/data', is_train=True, grnd_scale_distri=train_scale_distri, grnd_rotation_distri=train_rotation_distri) train_loader = torch.utils.data.DataLoader(dataset, batch_size=opt.bs, shuffle=True, num_workers=opt.nThreads, pin_memory=True) print('collecting validation scale and rotation distributions ...\n') val_scale_distri = read_grnd_distri_from_txt(val_scale_path) val_rotation_distri = read_grnd_distri_from_txt(val_rotation_path) dataset = MPII('dataset/mpii-hr-lsp-normalizer.json', '/bigdata1/zt53/data', is_train=False, grnd_scale_distri=val_scale_distri, grnd_rotation_distri=val_rotation_distri) val_loader = torch.utils.data.DataLoader(dataset, batch_size=opt.bs, shuffle=False, num_workers=opt.nThreads, pin_memory=True) agent = model.create_asn(chan_in=256, chan_out=256, scale_num=len(dataset.scale_means), rotation_num=len(dataset.rotation_means), is_aug=True) agent = torch.nn.DataParallel(agent).cuda() optimizer = torch.optim.RMSprop(agent.parameters(), lr=opt.lr, alpha=0.99, eps=1e-8, momentum=0, weight_decay=0) # optimizer = torch.optim.Adam(agent.parameters(), lr=opt.agent_lr) if opt.load_prefix_sr == '': checkpoint_agent.save_prefix = sr_pretrain_dir + '/' else: checkpoint_agent.save_prefix = sr_pretrain_dir + '/' + opt.load_prefix_sr checkpoint_agent.load_prefix = checkpoint_agent.save_prefix[0:-1] checkpoint_agent.load_checkpoint(agent, optimizer, train_history) # adjust_lr(optimizer, opt.lr) # lost_joint_count_path = os.path.join(opt.exp_dir, opt.exp_id, opt.asdn_dir, 'joint-count-finetune.txt') print('agent: ', type(optimizer), optimizer.param_groups[0]['lr']) if opt.dataset == 'mpii': num_classes = 16 hg = model.create_hg(num_stacks=2, num_modules=1, num_classes=num_classes, chan=256) hg = torch.nn.DataParallel(hg).cuda() if opt.load_prefix_pose == '': print('please input the checkpoint name of the pose model') exit() checkpoint_hg = Checkpoint() # checkpoint_hg.save_prefix = os.path.join(opt.exp_dir, opt.exp_id, opt.resume_prefix_pose) checkpoint_hg.load_prefix = os.path.join(opt.exp_dir, opt.exp_id, opt.load_prefix_pose)[0:-1] checkpoint_hg.load_checkpoint(hg) logger = Logger(sr_pretrain_dir + '/' + 'training-summary.txt', title='training-summary') logger.set_names(['Epoch', 'LR', 'Train Loss', 'Val Loss']) """training and validation""" start_epoch = 0 if opt.load_prefix_sr != '': start_epoch = train_history.epoch[-1]['epoch'] + 1 for epoch in range(start_epoch, opt.nEpochs): # train for one epoch train_loss = train(train_loader, hg, agent, optimizer, epoch, visualizer, opt) val_loss = validate(val_loader, hg, agent, epoch, visualizer, opt) # update training history e = OrderedDict([('epoch', epoch)]) lr = OrderedDict([('lr', optimizer.param_groups[0]['lr'])]) loss = OrderedDict([('train_loss', train_loss), ('val_loss', val_loss)]) # pckh = OrderedDict( [('val_pckh', val_pckh)] ) train_history.update(e, lr, loss) # print(train_history.lr[-1]['lr']) checkpoint_agent.save_checkpoint(agent, optimizer, train_history, is_asn=True) visualizer.plot_train_history(train_history, 'sr') logger.append( [epoch, optimizer.param_groups[0]['lr'], train_loss, val_loss]) logger.close()