def forward_backward_analysis(nn: Network=None, fraction=1): if nn is None: nn = Network() print("Loading Data") images, labels = get_test_set() shuffle_sets(images, labels) pool_size = 4 data_size = images.shape[0]//fraction splits = int(np.sqrt(np.maximum(10-fraction, 0))+1) size = data_size//(pool_size*splits) pool = Pool(pool_size) data = [(nn, images[i:i+size], labels[i:i+size]) for i in range(0, data_size, size)] print("Calculating Node Statistics", end='', flush=True) result = pool.starmap(_combine_fb_analysis, data) print() weights_sum = result[0][0] weights_max = result[0][1] biases = result[0][2] correct = result[0][3] weights_num = result[0][4] for i in range(1, len(result)): for j, _ in enumerate(weights_sum): weights_sum[j] += result[i][0][j] weights_max[j] = np.maximum(weights_max[j], result[i][0][j]) biases += result[i][2] correct += result[i][3] for j, _ in enumerate(weights_num): weights_num[j] += result[i][4][j] for i in range(len(weights_sum)): weights_sum[i] /= data_size biases /= data_size weights = [np.maximum(a, b) for a, b in zip(weights_max, weights_sum)] return weights, biases, float(correct)/float(size*pool_size*splits), weights_num
def main(): global opt opt = parser.parse_args() opt.gpuids = list(map(int, opt.gpuids)) print(opt) if opt.cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") cudnn.benchmark = True train_set = get_training_set(opt.upscale_factor, opt.add_noise, opt.noise_std) validation_set = get_validation_set(opt.upscale_factor) test_set = get_test_set(opt.upscale_factor) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batch_size, shuffle=True) validating_data_loader = DataLoader(dataset=validation_set, num_workers=opt.threads, batch_size=opt.test_batch_size, shuffle=False) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.test_batch_size, shuffle=False) model = SRCNN() criterion = nn.MSELoss() if opt.cuda: torch.cuda.set_device(opt.gpuids[0]) with torch.cuda.device(opt.gpuids[0]): model = model.cuda() criterion = criterion.cuda() model = nn.DataParallel(model, device_ids=opt.gpuids, output_device=opt.gpuids[0]) optimizer = optim.Adam(model.parameters(), lr=opt.lr) if opt.test: model_name = join("model", opt.model) model = torch.load(model_name) model = nn.DataParallel(model, device_ids=opt.gpuids, output_device=opt.gpuids[0]) start_time = time.time() test(model, criterion, testing_data_loader) elapsed_time = time.time() - start_time print("===> average {:.2f} image/sec for processing".format( 100.0 / elapsed_time)) return for epoch in range(1, opt.epochs + 1): train(model, criterion, epoch, optimizer, training_data_loader) validate(model, criterion, validating_data_loader) if epoch % 10 == 0: checkpoint(model, epoch)
def load_dataset(self, dataset, is_train=True): if self.num_channels == 1: is_gray = True else: is_gray = False if is_train: print('Loading train datasets...') train_set = get_training_set(self.data_dir, dataset, self.crop_size, self.scale_factor, is_gray=is_gray) return DataLoader(dataset=train_set, num_workers=self.num_threads, batch_size=self.batch_size, shuffle=True) else: print('Loading test datasets...') test_set = get_test_set(self.data_dir, dataset, self.scale_factor, is_gray=is_gray) return DataLoader(dataset=test_set, num_workers=self.num_threads, batch_size=self.test_batch_size, shuffle=False)
def test_dataloader(self): test_set = get_test_set(self.upscaleFactor) testing_data_loader = DataLoader(dataset=test_set, num_workers=4, batch_size=10, shuffle=False) return testing_data_loader
def main(): sys.stdout = Logger( os.path.join(opt.save_test_log, 'test_' + systime + '.txt')) if not torch.cuda.is_available(): raise Exception('No Gpu found, please run with gpu') else: use_gpu = torch.cuda.is_available() if use_gpu: cudnn.benchmark = False torch.cuda.manual_seed(opt.seed) pin_memory = True if use_gpu else False n_c = 128 n_b = 5 rrn = RRN(opt.scale, n_c, n_b) # initial filter generate network print(rrn) print("Model size: {:.5f}M".format( sum(p.numel() for p in rrn.parameters()) * 4 / 1048576)) print('===> {}L model has been initialized'.format(n_b)) rrn = torch.nn.DataParallel(rrn, device_ids=gpus_list) print('===> load pretrained model') if os.path.isfile(opt.pretrain): rrn.load_state_dict( torch.load(opt.pretrain, map_location=lambda storage, loc: storage)) print('===> pretrained model is load') else: raise Exception('pretrain model is not exists') if use_gpu: rrn = rrn.cuda(gpus_list[0]) print('===> Loading test Datasets') PSNR_avg = 0 SSIM_avg = 0 count = 0 scene_list = os.listdir(opt.test_dir) #scene_list = ['calendar','city','foliage','walk'] # Vid4 #scene_list = ['archpeople','archwall','auditorium','band','caffe','camera','lake','clap','photography','polyflow'] # UDM10 #scene_list = ['car05_001', 'hdclub_003_001', 'hitachi_isee5_001', 'hk004_001', 'HKVTG_004', 'jvc_009_001', 'NYVTG_006', 'PRVTG_012', 'RMVTG_011', 'veni3_011', 'veni5_015'] # SPMCS for scene_name in scene_list: test_set = get_test_set(opt.test_dir, opt.scale, scene_name) test_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testbatchsize, shuffle=False, pin_memory=pin_memory, drop_last=False) print('===> DataLoading Finished') PSNR, SSIM = test(test_loader, rrn, opt.scale, scene_name, n_c) PSNR_avg += PSNR SSIM_avg += SSIM count += 1 PSNR_avg = PSNR_avg / len(scene_list) SSIM_avg = SSIM_avg / len(scene_list) print('==> Average PSNR = {:.6f}'.format(PSNR_avg)) print('==> Average SSIM = {:.6f}'.format(SSIM_avg))
def main(): global opt, model opt = parser.parse_args() print(opt) cudnn.benchmark = True print('Loading datasets') train_set = get_training_set() test_set = get_test_set() training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) model = Net() criterion = L1_Charbonnier_loss() print("GPU will be used if found") if torch.cuda.is_available(): model = model.cuda() criterion = criterion.cuda() # optionally resume from a checkpoint if opt.resume: if os.path.isfile(opt.resume): print("loading checkpoint '{}'".format(opt.resume)) checkpoint = torch.load(opt.resume) opt.start_epoch = checkpoint["epoch"] + 1 model.load_state_dict(checkpoint["model"].state_dict()) else: print("no checkpoint found at '{}'".format(opt.resume)) # optionally copy weights from a checkpoint if opt.pretrained: if os.path.isfile(opt.pretrained): print("loading model '{}'".format(opt.pretrained)) weights = torch.load(opt.pretrained) model.load_state_dict(weights['model'].state_dict()) else: print("no model found at '{}'".format(opt.pretrained)) optimizer = optim.Adam(model.parameters(), lr=opt.lr) print("Training") for epoch in range(opt.start_epoch, opt.nEpochs + 1): train(training_data_loader, optimizer, model, criterion, epoch) test() if epoch % 100: save_checkpoint(model, epoch)
def main(): parser = argparse.ArgumentParser(description='PyTorch LapSRN') parser.add_argument('--batchSize', type=int, default=20, help='training batch size') parser.add_argument('--train_dir', type=str, default=None, help='path to train dir') parser.add_argument('--testBatchSize', type=int, default=10, help='testing batch size') parser.add_argument('--nEpochs', type=int, default=10, help='number of epochs to train for') parser.add_argument('--pretrained_vgg', type=int, default=0, help='vgg for perceptual loss') parser.add_argument('--checkpoint', type=str, default=MODEL_DIR, help='Path to checkpoint') parser.add_argument('--lr', type=float, default=1e-2, help='Learning Rate. Default=1e2') parser.add_argument('--threads', type=int, default=4, help='number of threads for data loader to use') parser.add_argument('--seed', type=int, default=123, help='random seed to use. Default=123') parser.add_argument('--loss_type', type=str, default='pl', help='MSE or PerceptualLoss loss (mse, pl)') opt = parser.parse_args() if not torch.cuda.is_available(): raise Exception("No GPU found") torch.manual_seed(opt.seed) random.seed(opt.seed) torch.cuda.manual_seed(opt.seed) print('Loading datasets') train_set = get_training_set(opt.train_dir) test_set = get_test_set() training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print 'Building model' model = LasSRN() model = model.cuda() if opt.loss_type == 'pl': print 'Loading VGG-style (y-channel) for perceptual loss' criterion = PerceptualLoss(*create_discriptor_net(layers=['relu_4', 'relu_6', 'relu_8', 'relu_10'], pretrained=opt.pretrained_vgg)) else: criterion = nn.MSELoss() criterion = criterion.cuda() lr = opt.lr optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-5) # TODO, to config print 'Starting learning', 'lr=', lr, 'batch_size=', opt.batchSize for epoch in range(1, opt.nEpochs + 1): print 'Epoch num=%d' % epoch train(epoch, model, criterion, optimizer, training_data_loader) test(model, torch.nn.MSELoss().cuda(), testing_data_loader) if epoch: lr = lr/2 print 'Setting learing rate to %f' % (lr, ) for param_group in optimizer.param_groups: param_group['lr'] = lr checkpoint(opt, epoch, model)
def main(): #writer = SummaryWriter() sys.stdout = Logger( os.path.join(opt.save_test_log, 'test_' + systime + '.txt')) if not torch.cuda.is_available(): raise Exception('No Gpu found, please run with gpu') else: use_gpu = torch.cuda.is_available() if use_gpu: cudnn.benchmark = False torch.cuda.manual_seed(opt.seed) pin_memory = True if use_gpu else False # Selecting network rsdn = RSDN9_128(4) # initial filter generate network print(rsdn) print("Model size: {:.5f}M".format( sum(p.numel() for p in rsdn.parameters()) * 4 / 1048576)) rsdn = torch.nn.DataParallel(rsdn, device_ids=gpus_list) print('===> load pretrained model') if os.path.isfile(opt.pretrain): rsdn.load_state_dict( torch.load(opt.pretrain, map_location=lambda storage, loc: storage)) print('===> pretrained model is load') else: raise Exception('pretrain model is not exists') if use_gpu: rsdn = rsdn.cuda(gpus_list[0]) print('===> Loading test Datasets') PSNR_avg = 0 SSIM_avg = 0 count = 0 out = [] test_list = ['foliage_r.txt', 'walk_r.txt', 'city_r.txt', 'calendar_r.txt'] for test_name in test_list: test_set = get_test_set(opt.test_dir, opt.scale, test_name.split('.')[0]) test_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testbatchsize, shuffle=False, pin_memory=pin_memory, drop_last=False) print('===> DataLoading Finished') PSNR, SSIM, out = test(test_loader, rsdn, test_name.split('.')[0], out) PSNR_avg += PSNR SSIM_avg += SSIM count += 1 PSNR_avg = PSNR_avg / len(test_list) SSIM_avg = SSIM_avg / len(test_list) print('==> Average PSNR = {:.6f}'.format(PSNR_avg)) print('==> Average SSIM = {:.6f}'.format(SSIM_avg))
def main(): sys.stdout = Logger( os.path.join(opt.save_test_log, 'test_' + systime + '.txt')) if not torch.cuda.is_available(): raise Exception('No Gpu found, please run with gpu') else: use_gpu = torch.cuda.is_available() if use_gpu: cudnn.benchmark = False torch.cuda.manual_seed(opt.seed) pin_memory = True if use_gpu else False # Selecting network net = TGA(4) # initial filter generate network print("Model size: {:.5f}M".format( sum(p.numel() for p in net.parameters()) * 4 / 1048576)) print('===> {} model has been initialized'.format(opt.model_name)) filter_net = torch.nn.DataParallel(net, device_ids=gpus_list) print('===> load pretrained model') if os.path.isfile(opt.pretrain): filter_net.load_state_dict( torch.load(opt.pretrain, map_location=lambda storage, loc: storage)) print('===> pretrained model is load') else: raise Exception('pretrain model is not exists') if use_gpu: net = net.cuda(gpus_list[0]) print('===> Loading test Datasets') PSNR_avg = 0 SSIM_avg = 0 test_list = ['calendar.txt', 'city.txt', 'foliage.txt', 'walk.txt'] for test_name in test_list: opt.file_test_list = test_name test_set = get_test_set(opt.test_dir, opt.file_test_list, opt.scale, opt.nFrames) test_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testbatchsize, shuffle=False, pin_memory=pin_memory, drop_last=False) print('===> DataLoading Finished') PSNR, SSIM = test(test_loader, net, test_name.split('.')[0]) PSNR_avg += PSNR SSIM_avg += SSIM PSNR_avg = PSNR_avg / len(test_list) SSIM_avg = SSIM_avg / len(test_list) print('==> Average PSNR = {:.6f}'.format(PSNR_avg)) print('==> Average SSIM = {:.6f}'.format(SSIM_avg))
def export_predictions(): """ Export the predictions and data as a csv """ img, lab = get_test_set() model = tf.estimator.Estimator(model_fn, FOLDER) inp = tf.estimator.inputs.numpy_input_fn(dict(img=img), lab, 1000, 1, False, 10000) with open("export.csv", "w") as file: for i, p in enumerate(model.predict(inp)): file.write("%d, %d" % (lab[i], p["prediction"])) for k in p["probabilities"]: file.write(", %f" % k) for k in np.nditer(img[i, :, :, :]): file.write(", %f" % k) file.write("\n")
def main(): # =========================================================== # Set train dataset & test dataset # =========================================================== print('===> Loading datasets') train_set = get_training_set(args.upscale_factor) test_set = get_test_set(args.upscale_factor) training_data_loader = DataLoader(dataset=train_set, batch_size=args.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, batch_size=args.testBatchSize, shuffle=False) model = SRCNNTrainer(args, training_data_loader, testing_data_loader) model.run()
def main(not_parsed_args): logging.info('Loading evaluation dataset...') test_dataset = get_test_set(FLAGS.input_dir, None, FLAGS.frames, 4, 'filelist.txt', True, FLAGS.future_frame) logging.info('done. size %d' % len(test_dataset)) logging.info('Loading model...') model = RBPN() model.load_weights(FLAGS.model) logging.info('Test start...') for s in range(len(test_dataset)): logging.info('Step %d' % s) x, _, filenames = test_dataset.batch(withName=True) y = model.predict(x) y *= 255.0 Image.fromarray(y[0].astype(np.uint8), mode='RGB').save( os.path.join(FLAGS.output_dir, filenames[0])) logging.info('Test finished')
def evaluate(train=False, epochs=None): """ Evaluate the fitness of this model Arguments: train: train the model before evaluating epochs: number of epochs to train (None means infinite) """ tf.logging.set_verbosity(tf.logging.INFO) model = tf.estimator.Estimator(model_fn, FOLDER) if train: img, lab = get_training_set() input_fn = tf.estimator.inputs.numpy_input_fn(dict(img=img), lab, 128, epochs, True, 6000) model.train(input_fn) img, lab = get_test_set() return model.evaluate( tf.estimator.inputs.numpy_input_fn(dict(img=img), lab, 1000, 1, True, 10000))
def __init__(self, module_path, hyper_params, use_cuda, mission=1): self.dataset = get_test_set(mission=mission) print("test number:", len(self.dataset)) self.hyper_params = hyper_params self.data_loader = DataLoader( dataset=self.dataset, num_workers=self.hyper_params["threads"], batch_size=self.hyper_params["batch_size"], shuffle=False) self.resnet = get_network() self.resnet.load_state_dict(torch.load(module_path)) if use_cuda: self.resnet = self.resnet.cuda() self.v = Validator(resnet=self.resnet, hyper_params=hyper_params, use_cuda=use_cuda, data_loader=self.data_loader)
def evaluate(self, set_size=1.0, print_result=True): """ Evaluate the network """ img, lab = get_test_set() img.shape = img.shape[0], np.prod(img.shape[1:]) correct = 0 if set_size < 1.0 and set_size >= 0.0: shuffle_sets(img, lab) size = int(lab.shape[0] * set_size) else: size = lab.shape[0] for i in range(size): prev = img[i] for layer in self.layers: prev, _ = layer.forward(prev, 0) if np.argmax(prev) == lab[i]: correct += 1 if print_result: print('Evaluation:', correct, '/', size, 'correct') return float(correct) / float(size)
def __init__(self, model_path, save_dir, hyper_params, use_cuda, mission=1): self.dataset = get_test_set(mission=mission) self.hyper_params = hyper_params self.data_loader = DataLoader( dataset=self.dataset, num_workers=self.hyper_params["threads"], batch_size=self.hyper_params["batch_size"], shuffle=False) self.save_dir = save_dir self.model_path = model_path self.use_cuda = use_cuda self.mission = mission self.resnet = get_network(mission=mission) self.resnet.load_state_dict(torch.load(model_path)) if use_cuda: self.resnet = self.resnet.cuda() pass
default=123, help='random seed to use. Default=123') opt = parser.parse_args() print(opt) if opt.cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") torch.manual_seed(opt.seed) device = torch.device("cuda" if opt.cuda else "cpu") print('===> Loading datasets') train_set = get_training_set(opt.upscale_factor) test_set = get_test_set(opt.upscale_factor) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print('===> Building model') model = Net(upscale_factor=opt.upscale_factor).to(device) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=opt.lr)
def main(): print(opt) cuda = opt.cuda if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") opt.seed = random.randint(1, 10000) torch.manual_seed(opt.seed) if cuda: torch.cuda.manual_seed(opt.seed) cudnn.benchmark = True print('===> Loading datasets') train_set = get_training_set(opt.dataset) test_set = get_test_set(opt.dataset) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) test_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print("===> Building model") if (opt.net == 'resnet'): model = ResNet() else: model = TFNet() criterion = nn.L1Loss() print("===> Setting GPU") if cuda: model = model.cuda() criterion = criterion.cuda() # optionally resume from a checkpoint if opt.resume: if os.path.isfile(opt.resume): print("=> loading checkpoint '{}'".format(opt.resume)) checkpoint = torch.load(opt.resume) opt.start_epoch = checkpoint["epoch"] + 1 model.load_state_dict(checkpoint["model"].state_dict()) else: print("=> no checkpoint found at '{}'".format(opt.resume)) # optionally copy weights from a checkpoint if opt.pretrained: if os.path.isfile(opt.pretrained): print("=> loading model '{}'".format(opt.pretrained)) weights = torch.load(opt.pretrained) model.load_state_dict(weights['model'].state_dict()) else: print("=> no model found at '{}'".format(opt.pretrained)) print("===> Setting Optimizer") optimizer = optim.Adam(model.parameters(), lr=opt.lr) print("===> Training") t = time.strftime("%Y%m%d%H%M") train_log = open( os.path.join(opt.log, "%s_%s_train.log") % (opt.net, t), "w") test_log = open( os.path.join(opt.log, "%s_%s_test.log") % (opt.net, t), "w") for epoch in range(opt.start_epoch, opt.nEpochs + 1): train(training_data_loader, optimizer, model, criterion, epoch, train_log) if epoch % 10 == 0: test(test_data_loader, model, criterion, epoch, test_log) save_checkpoint(model, epoch, t) train_log.close() test_log.close()
opt = parser.parse_args() gpus_list = range(opt.gpus) print(opt) cuda = opt.gpu_mode if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") torch.manual_seed(opt.seed) if cuda: torch.cuda.manual_seed(opt.seed) print('===> Loading datasets') test_set = get_test_set(opt.data_dir, opt.upscale_factor, opt.file_list) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print('===> Building model ', opt.model_type) if opt.model_type == 'FBPNSR_RBPN_V1_REF': model = FBPNSR_RBPN_V1_REF(base_filter=256, feat=64, num_stages=3, n_resblock=5, scale_factor=opt.upscale_factor) elif opt.model_type == 'FBPNSR_RBPN_V2_REF': model = FBPNSR_RBPN_V2_REF(base_filter=256, feat=64,
print(opt) if opt.cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") cudnn.benchmark = True torch.manual_seed(opt.seed) if opt.cuda: torch.cuda.manual_seed(opt.seed) print('===> Loading datasets') root_path = "dataset/" train_set = get_training_set(root_path + opt.dataset) test_set = get_test_set(root_path + opt.dataset) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print('===> Building model') netG = G(opt.input_nc, opt.output_nc, opt.ngf) pytorch_total_params = sum(p.numel() for p in netG.parameters() if p.requires_grad) print ("\nTrainable parameters", pytorch_total_params) #netG.apply(weights_init) netD = D(opt.input_nc, opt.output_nc, opt.ndf) #netD.apply(weights_init) criterion = nn.BCELoss() criterion_l1 = nn.L1Loss() criterion_mse = nn.MSELoss()
default=123, help='Random seed to use. Default=123') opt = parser.parse_args() print(opt) if opt.cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") torch.manual_seed(opt.seed) device = torch.device("cuda" if opt.cuda else "cpu") full_size = opt.full_size print('===> Loading datasets') train_set = get_training_set(opt.upscale_factor, opt.full_size) test_set = get_test_set(opt.upscale_factor, opt.full_size) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=True) print('===> Building model') model = RNet(upscale_factor=opt.upscale_factor, full_size=opt.full_size) model.to(device) criterion = nn.MSELoss() #Three optimizers, one for each output optimizerLow = optim.Adam(model.parameters(), lr=opt.lr)
print(opt) if opt.cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") cudnn.benchmark = True # 加速 torch.manual_seed(opt.seed) if opt.cuda: torch.cuda.manual_seed(opt.seed) print('===> Loading datasets') root_path = "dataset/" train_set = get_training_set(root_path + opt.dataset) #objectlize test_set = get_test_set(root_path + opt.dataset) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print('===> Building model') netG = define_G(opt.input_nc, opt.output_nc, opt.ngf, 'batch', False, [0]) netD = define_D(opt.input_nc + opt.output_nc, opt.ndf, 'batch', False, [0]) #opt.input_nc + opt.output_nc, 因为cat( ,1) criterionGAN = GANLoss()
from __future__ import print_function from math import log10 import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from model import ConvSubPixel from data import get_training_set, get_test_set import os torch.manual_seed(123) print('Loading datasets...') train_set = get_training_set(8) test_set = get_test_set(8) training_data_loader = DataLoader(dataset=train_set, num_workers=4, batch_size=32, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=4, batch_size=32, shuffle=False) print('Building model...') model = ConvSubPixel(upscale_factor=8) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=1e-2)
parser.add_argument('--scale', type=int, default=4, help='scale output size /input size') opt = parser.parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = "0" print(opt) cuda = opt.gpu_mode if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") torch.manual_seed(opt.seed) if cuda: torch.cuda.manual_seed(opt.seed) print('===> Loading datasets') test_set = get_test_set(opt.data_dir, opt.file_list, 5) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print('===> Building model') model = DAVSR_1(opt.nChannel, opt.nFeat, opt.scale) device = torch.device("cuda") model.load_state_dict(torch.load(opt.model)) model.to(device) print('Pre-trained SR model is loaded.') if cuda:
def train(training_opt: TrainingOptions, device_opt: DeviceOptions) -> Net: if device_opt.cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") torch.manual_seed(training_opt.seed) device = torch.device("cuda" if device_opt.cuda else "cpu") print('===> Loading datasets') train_set = get_train_set(training_opt.upscale_factor) test_set = get_test_set(training_opt.upscale_factor) training_data_loader = DataLoader(dataset=train_set, num_workers=device_opt.threads, batch_size=device_opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=device_opt.threads, batch_size=device_opt.testBatchSize, shuffle=False) print('===> Building model') model = Net(upscale_factor=training_opt.upscale_factor).to(device) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=training_opt.lr) training_status = st.empty() test_status = st.empty() def train(epoch: int): training_status.code('Training... (epoch=%s)' % epoch) epoch_loss = 0 for iteration, batch in enumerate(training_data_loader, 1): input, target = batch[0].to(device), batch[1].to(device) optimizer.zero_grad() loss = criterion(model(input), target) epoch_loss += loss.item() loss.backward() optimizer.step() training_status.code( "===> Epoch[{}/{}]({}/{}): Loss: {:.4f}".format( epoch, training_opt.nEpochs, iteration, len(training_data_loader), loss.item())) training_status.code( "===> Epoch {} Complete: Avg. Loss: {:.4f}".format( epoch, epoch_loss / len(training_data_loader))) def test(): avg_psnr = 0 with torch.no_grad(): for batch in testing_data_loader: input, target = batch[0].to(device), batch[1].to(device) prediction = model(input) mse = criterion(prediction, target) psnr = 10 * log10(1 / mse.item()) avg_psnr += psnr test_status.code("===> Avg. PSNR: {:.4f} dB".format( avg_psnr / len(testing_data_loader))) for epoch in range(1, training_opt.nEpochs + 1): train(epoch) # test() return model
gpus_list = range(args.gpus) print(args) cuda = args.gpu_mode if cuda: print("Using GPU mode") if not torch.cuda.is_available(): raise Exception("No GPU found, please run without --gpu_mode") torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed(args.seed) print('==> Loading datasets') test_set = get_test_set(args.data_dir, args.nFrames, args.upscale_factor, args.file_list, args.other_dataset, args.future_frame) testing_data_loader = DataLoader(dataset=test_set, num_workers=args.threads, batch_size=args.testBatchSize, shuffle=False) print('==> Building model ', args.model_type) if args.model_type == 'RBPN': model = RBPN(num_channels=3, base_filter=256, feat=64, num_stages=3, n_resblock=5, nFrames=args.nFrames, scale_factor=args.upscale_factor)
print(opt) if opt.cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") cudnn.benchmark = True torch.manual_seed(opt.seed) if opt.cuda: torch.cuda.manual_seed(opt.seed) print('===> Loading datasets') root_path = "dataset/" train_set = get_training_set(root_path + opt.dataset, opt.direction) test_set = get_test_set(root_path + opt.dataset, opt.direction) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batch_size, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.test_batch_size, shuffle=False) device = torch.device("cuda:0" if opt.cuda else "cpu") print('===> Building models') net_g = define_G(opt.input_nc, opt.output_nc, opt.ngf,
parser.add_argument('--threads', type=int, default=4, help='number of threads for data loader to use') parser.add_argument('--seed', type=int, default=123, help='random seed to use. Default=123') opt = parser.parse_args() print(opt) if opt.cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") torch.manual_seed(opt.seed) device = torch.device("cuda" if opt.cuda else "cpu") print('===> Loading datasets') train_set = get_training_set(opt.upscale_factor) test_set = get_test_set(opt.upscale_factor) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print('===> Building model') model = Net(upscale_factor=opt.upscale_factor).to(device) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=opt.lr) def train(epoch): epoch_loss = 0 for iteration, batch in enumerate(training_data_loader, 1): input, target = batch[0].to(device), batch[1].to(device)
#Uncomment threads while using a powerful gpu #threads = 4 seed = 123 if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run with cuda == False") else: device = 'cuda' torch.manual_seed(seed) #device = torch.device("cuda" if cuda else "cpu") print('===> Loading datasets') train_set = get_training_set(upscale_factor) test_set = get_test_set(upscale_factor) #To use with threads #training_data_loader = DataLoader(dataset=train_set, num_workers=threads, batch_size=batchSize, shuffle=True) #testing_data_loader = DataLoader(dataset=test_set, num_workers=threads, batch_size=testBatchSize, shuffle=False) training_data_loader = DataLoader(dataset=train_set, batch_size=batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, batch_size=testBatchSize, shuffle=False) print('===> Building model') model = Net(num_channels=1, upscale_factor=upscale_factor).to(device) criterion = nn.MSELoss()
return x.cpu().data.numpy() cuda = opt.cuda print(cuda) if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") torch.manual_seed(opt.seed) if cuda: torch.cuda.manual_seed(opt.seed) print('===> Loading datasets') train_set = get_training_set(opt.upscale_factor, opt.dataset) test_set = get_test_set(opt.upscale_factor, opt.dataset) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) num_features = opt.num_features filter_size = opt.filter_size batch_size = opt.batchSize shape = (84 / opt.upscale_factor, 84 / opt.upscale_factor) #H,W inp_chans = opt.inp_chans nlayers = opt.nlayers
opt = parser.parse_args() gpus_list = range(opt.gpus) print(opt) cuda = opt.gpu_mode if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") torch.manual_seed(opt.seed) if cuda: torch.cuda.manual_seed(opt.seed) print('===> Loading datasets') test_set = get_test_set(opt.data_dir, opt.nFrames, opt.upscale_factor, opt.file_list, opt.other_dataset, opt.future_frame) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print('===> Building model ', opt.model_type) if opt.model_type == 'RBPN': model = RBPN(num_channels=3, base_filter=256, feat=64, num_stages=3, n_resblock=5, nFrames=opt.nFrames, scale_factor=opt.upscale_factor)
print(opt) np.save("checkpoint/{}/parameters_train_{}".format(opt.dataset, strftime("%H%M", gmtime())), opt) if opt.cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") cudnn.benchmark = True torch.manual_seed(opt.seed) if opt.cuda: torch.cuda.manual_seed(opt.seed) print('===> Loading datasets') root_path = "dataset/" train_set = get_training_set(root_path + opt.dataset, mode=opt.mode, xy=opt.xy_size, img_size=(opt.img_w, opt.img_h), norm=opt.norm) test_set = get_test_set(root_path + opt.dataset, mode = opt.mode, xy=opt.xy_size, img_size=(opt.img_w, opt.img_h), norm=opt.norm) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print('===> Building model') netG = define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.xy_size, 'batch', False, [0]) netD = define_D(opt.input_nc + opt.output_nc, opt.ndf, 'batch', False, [0]) if opt.goon!=0: #Start mith on-the-last-iteration saved model netG = torch.load("checkpoint/cityscapes/netG_model_epoch_{}.pth".format(opt.goon)) netD = torch.load("checkpoint/cityscapes/netD_model_epoch_{}.pth".format(opt.goon)) criterionGAN = GANLoss() criterionL1 = nn.L1Loss()