def __init__(self, args, device, dataloader_sup, dataloader_unsup, dataloader_test, n_classes): if args.dataset == 'mnist': from models import Autoencoder, VaDE self.autoencoder = Autoencoder().to(device) self.autoencoder.apply(weights_init_normal) self.VaDE = VaDE().to(device) elif args.dataset == 'webcam': from models_office import Autoencoder, VaDE, feature_extractor self.autoencoder = Autoencoder().to(device) checkpoint = torch.load('weights/imagenet_params.pth.tar', map_location=device) self.autoencoder.load_state_dict(checkpoint['state_dict'], strict=False) checkpoint = torch.load('weights/feature_extractor_params.pth.tar', map_location=device) self.feature_extractor = feature_extractor().to(device) self.feature_extractor.load_state_dict(checkpoint['state_dict']) self.freeze_extractor() self.VaDE = VaDE().to(device) self.dataloader_sup = dataloader_sup self.dataloader_unsup = dataloader_unsup self.dataloader_test = dataloader_test self.device = device self.args = args self.n_classes = n_classes
def get_model(args): ''' define model ''' use_unet = True if args.model == 'UNet' else False model = Autoencoder(use_unet) print('---Model Information---') print('Net:', model) print('Use GPU:', args.use_cuda) return model.to(args.device)
def train_autoencoder(epoch_size=250): start_time = datetime.datetime.now() print(f"train_encoder start time: {str(start_time)}") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print('Using device:', device) train_dataset = EncodedDeepfakeDataset(train_folders, encoder=None, n_frames=1, device=device) model = Autoencoder() model = model.to(device) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters()) num_epochs = 10 batch_size = 1 dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) for epoch in range(num_epochs): epoch_start_time = datetime.datetime.now() epoch_loss = 0 for i, batch in enumerate(dataloader): if i * batch_size >= epoch_size: break data, _, _ = batch data = data.to(device) data = data.squeeze(0) optimizer.zero_grad() output = model(data) loss = criterion(output, data) epoch_loss += loss.item() loss.backward() optimizer.step() epoch_end_time = datetime.datetime.now() exec_time = epoch_end_time - epoch_start_time print( f'epoch: {epoch}, loss: {epoch_loss/(epoch_size/batch_size)}, executed in: {str(exec_time)}' ) end_time = datetime.datetime.now() torch.save(model.state_dict(), f'autoencoder_{end_time.strftime("H%HM%MS%S_%m-%d-%y")}.pt') exec_time = end_time - start_time print( f"train_encoder executed in: {str(exec_time)}, end time: {str(end_time)}" ) if device.type == 'cuda': print(torch.cuda.get_device_name(0)) print('Memory Usage:') print('Allocated:', round(torch.cuda.memory_allocated(0) / 1024**3, 1), 'GB') print('Cached: ', round(torch.cuda.memory_cached(0) / 1024**3, 1), 'GB') return (loss.item(), exec_time)
def convert_face(croped_face): resized_face = cv2.resize(croped_face, (256, 256)) normalized_face = resized_face / 255.0 #normalized_face = normalized_face.reshape(1, normalized_face.shape[0], normalized_face.shape[1], normalized_face.shape[2]) warped_img, _ = random_warp(normalized_face) batch_warped_img = np.expand_dims(warped_img, axis=0) batch_warped_img = toTensor(batch_warped_img) batch_warped_img = batch_warped_img.to(device).float() #print(batch_warped_img.shape, batch_warped_img) model = Autoencoder().to(device) checkpoint = torch.load('./checkpoint/autoencoder.t7') model.load_state_dict(checkpoint['state']) converted_face = model(batch_warped_img, 'B') return converted_face
def __init__(self, cnf): # type: (Conf) -> None self.cnf = cnf # init code predictor self.code_predictor = CodePredictor() self.code_predictor = self.code_predictor.to(cnf.device) # init volumetric heatmap autoencoder self.autoencoder = Autoencoder() self.autoencoder.eval() self.autoencoder.requires_grad(False) self.autoencoder = self.autoencoder.to(cnf.device) # init optimizer self.optimizer = optim.Adam(params=self.code_predictor.parameters(), lr=cnf.lr) # init dataset(s) training_set = JTATrainingSet(cnf) test_set = JTAValidationSet(cnf) # init train/test loader self.train_loader = DataLoader(training_set, cnf.batch_size, num_workers=cnf.n_workers, shuffle=True) self.test_loader = DataLoader(test_set, batch_size=1, num_workers=cnf.n_workers, shuffle=False) # init logging stuffs self.log_path = cnf.exp_log_path print(f'tensorboard --logdir={cnf.project_log_path.abspath()}\n') self.sw = SummaryWriter(self.log_path) self.log_freq = len(self.train_loader) # starting values values self.epoch = 0 self.best_test_f1 = None # possibly load checkpoint self.load_ck()
class ESRDataset(Dataset): def __init__(self, dirpath, modelpath): self.cage = [ x.path for x in os.scandir(os.path.join(dirpath, 'cage')) if x.name.endswith('jpg') or x.name.endswith('.png') ] self.trump = [ x.path for x in os.scandir(os.path.join(dirpath, 'trump')) if x.name.endswith('jpg') or x.name.endswith('.png') ] self.real_dir = self.cage + self.trump # auto encoder initialization, generate low quality image checkpoint = torch.load(modelpath, map_location='cpu') self.model = Autoencoder() self.model.load_state_dict(checkpoint['state']) self.random_transform_args = { 'rotation_range': 10, 'zoom_range': 0.05, 'shift_range': 0.05, 'random_flip': 0.4, } def __getitem__(self, index, tp=True): # return 256*256 image cls = 'A' if index < len(self.cage) else 'B' img = cv2.resize(cv2.imread(self.real_dir[index]), (256, 256)) / 255.0 img, target_img = random_warp( random_transform(img, **self.random_transform_args)) img = self.model( Variable(torch.unsqueeze(self.toTensor(img).float(), 0)), cls) img = torch.squeeze(img, 0) target_img = self.toTensor(target_img) target_img = target_img.float() return {'LQ': img, 'GT': target_img} def __len__(self): return len(self.real_dir) def toTensor(self, img): img = torch.from_numpy(img.transpose((2, 0, 1))) return img
def main(): data = input_data.read_data_sets('/data/mnist/', one_hot=True) with tf.Session() as sess: print("Autoencoder Training Started\n\n") ae_config = AEConfig() ae_model = Autoencoder(config) ae_trainer = Trainer(ae_config, ae_model, data, sess) ae_trainer.train() print("Autoencoder Training Ended\n\n")
def __init__(self, save_path=SAVE_PATH, nb_image_to_gen=NB_IMAGE_TO_GENERATE, code_size=CODE_SIZE, image_size=IMAGE_SIZE, model_complexity=COMPLEXITY, mean=WEIGHTS_MEAN, std=WEIGHTS_STD, learning_rate=LEARNING_RATE, image_channels=IMAGE_CHANNELS, batch_size=MINIBATCH_SIZE): self.batch_size = batch_size self.save_path = save_path self.nb_image_to_gen = nb_image_to_gen self.image_size = image_size self.image_channels = image_channels # Device (cpu or gpu) self.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") # Model input_size = image_size * image_size * image_channels # times 3 because of colors self.autoencoder = Autoencoder(input_size, model_complexity, mean, std, code_size).to(self.device) self.optimiser = optim.Adam(self.autoencoder.parameters(), lr=learning_rate) self.losses = [] self.saved_code_input = torch.randn( (self.nb_image_to_gen * self.nb_image_to_gen, code_size)).to(self.device) # Create directory for the results if it doesn't already exists import os os.makedirs(self.save_path, exist_ok=True) os.makedirs(self.save_path + "encoded/", exist_ok=True) os.makedirs(self.save_path + "decoded/", exist_ok=True) os.makedirs(self.save_path + "saved_generated/", exist_ok=True)
def __init__(self, dirpath, modelpath): self.cage = [ x.path for x in os.scandir(os.path.join(dirpath, 'cage')) if x.name.endswith('jpg') or x.name.endswith('.png') ] self.trump = [ x.path for x in os.scandir(os.path.join(dirpath, 'trump')) if x.name.endswith('jpg') or x.name.endswith('.png') ] self.real_dir = self.cage + self.trump # auto encoder initialization, generate low quality image checkpoint = torch.load(modelpath, map_location='cpu') self.model = Autoencoder() self.model.load_state_dict(checkpoint['state']) self.random_transform_args = { 'rotation_range': 10, 'zoom_range': 0.05, 'shift_range': 0.05, 'random_flip': 0.4, }
def __init__(self): model = load_model() self.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") self.autoencoder = Autoencoder(model, self.device) self.predictor = Predictor(model, self.device) self.buffer_autoencoder_f = [] self.buffer_autoencoder_n = [] self.lim_autoencoder = 4000 self.threshold_autoencoder = 1 self.buffer_predictor_f = [] self.buffer_predictor_n = [] self.lim_predictor = 4000 self.threshold_predictor = 1 self.alpha = 0.95
def create_model(config, model_name, num_classes=1): section_training = config['TRAINING'] RGB = section_training.getboolean('RGB') if 'RGB' in section_training else None input_size = tuple(map(int, section_training.get('input_size').split('x'))) if 'input_size' in section_training else None section_hyper = config['HYPERPARAMS'] encode_factor = section_hyper.getint('encode_factor', None) base_channels = section_hyper.getint('base_channels', None) latent_dim = section_hyper.getint('latent_dim', 10) channel_increase_factor = section_hyper.getint('channel_increase_factor', 2) conv_blocks_per_decrease = section_hyper.getint('conv_blocks_per_decrease', 1) initial_upsample_size = section_hyper.getint('initial_upsample_size', 3) skip_connections = section_hyper.getboolean('skip_connections', False) num_classes = num_classes if config.getboolean('HYPERPARAMS', 'auxillary', fallback=False) else 0 """ if args.data == 'MNIST': input_size = (28,28) hidden_dim_size = (64, 2) encode_factor = 2 base_channels = 32 elif args.data == 'CelebA': input_size = (218, 178) hidden_dim_size = (64, 8) encode_factor = 4 base_channels = 32 else: raise NotImplementedError('Input size/ Hidden dimension for dataset not specified') """ ## Init model model = None if model_name in ['AE', 'VAE']: model = AE.Autoencoder(variational = model_name == 'VAE', final_size=input_size, encode_factor=encode_factor, RGB = RGB, base_channels=base_channels, channel_increase_factor=channel_increase_factor, conv_blocks_per_decrease=conv_blocks_per_decrease, encoding_dimension=latent_dim, initial_upsample_size=initial_upsample_size, skip_connections=skip_connections, n_classes=num_classes) elif model_name == 'VanillaGAN': latent_dim = config.getint('HYPERPARAMS', 'latent_dim', fallback=10) #print('Latent dim was set to {:d}'.format(latent_dim)) gen = GAN.VanillaGenerator(input_dim = latent_dim, num_classes=num_classes) disc = GAN.VanillaDiscriminator(n_classes=num_classes) model = (gen, disc) elif model_name == 'DCGAN': gen = GAN.DCGenerator(input_dim = latent_dim, num_classes=num_classes) disc = GAN.DCDiscriminator(n_classes=num_classes) model = (gen, disc) else: raise NotImplementedError('The model you specified is not implemented yet') return model
def test_autoencoder(): start_time = datetime.datetime.now() print(f"test_encoder start time: {str(start_time)}") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print('Using device:', device) with torch.no_grad(): test_folder = [ 'train/balanced_test', ] test_dataset = EncodedDeepfakeDataset(test_folder, encoder=None, n_frames=1, train=True) model = Autoencoder() model = model.to(device) criterion = nn.MSELoss() latest_ae = max(glob.glob('./autoencoder*.pt'), key=os.path.getctime) model.load_state_dict(torch.load(latest_ae, map_location=device)) num_vids = len(glob.glob(test_folder[0]+"/*.mp4")) sample_size = 400 sample = np.random.choice(a=num_vids, size=sample_size, replace=False) loss = 0 for ind in sample: data, _, _ = test_dataset[ind] data = data.to(device) out = model(data) loss += criterion(out, data).item() hidden = model.encode(data) end_time = datetime.datetime.now() exec_time = end_time - start_time print(f"executed in: {str(exec_time)}, finished {str(end_time)}") if device.type == 'cuda': print(torch.cuda.get_device_name(0)) print('Memory Usage:') print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB') print('Cached: ', round(torch.cuda.memory_cached(0)/1024**3,1), 'GB') return (loss / sample_size, exec_time, np.prod(list(hidden.size())))
else: device = torch.device("cpu") print("===> Using CPU to train") torch.manual_seed(args.seed) if args.cuda: torch.cuda.manual_seed(args.seed) print("===> Loaing datasets") images_A = get_image_paths("train/biden_face") images_B = get_image_paths("train/sacha_face") images_A = load_images(images_A) / 255.0 images_B = load_images(images_B) / 255.0 # images_A += images_B.mean(axis=(0, 1, 2)) - images_A.mean(axis=(0, 1, 2)) model = Autoencoder().to(device) print("===> Try resume from checkpoint") if os.path.isdir("checkpoint"): try: checkpoint = torch.load("./checkpoint/autoencoder.t7") model.load_state_dict(checkpoint["state"]) start_epoch = checkpoint["epoch"] print("===> Load last checkpoint data") except FileNotFoundError: print("Can't found autoencoder.t7") else: start_epoch = 0 print("===> Start from scratch") criterion = nn.L1Loss()
def results(cnf): # type: (Conf) -> None """ Shows a visual representation of the obtained results using the test set images as input """ # init Code Predictor code_predictor = CodePredictor() code_predictor.to(cnf.device) code_predictor.eval() code_predictor.requires_grad(False) code_predictor.load_w(f'log/{cnf.exp_name}/best.pth') # init Decoder autoencoder = Autoencoder(pretrained=True) autoencoder.to(cnf.device) autoencoder.eval() autoencoder.requires_grad(False) # init Hole Filler refiner = Refiner(pretrained=True) # refiner.to(cnf.device) refiner.eval() refiner.requires_grad(False) # init data loader ts = JTAValidationSet(cnf=cnf) loader = DataLoader(dataset=ts, batch_size=1, shuffle=False, num_workers=0) for step, sample in enumerate(loader): x, _, fx, fy, cx, cy, frame_path = sample x = x.to(cnf.device) fx, fy, cx, cy = fx.item(), fy.item(), cx.item(), cy.item() # image --> [code_predictor] --> code code_pred = code_predictor.forward(x).unsqueeze(0) # code --> [decode] --> hmap hmap_pred = autoencoder.decode(code_pred).squeeze() # hmap --> [local maxima search with cuda kernel] --> pseudo-3D coordinates pseudo3d_coords_pred = [] confidences = [] for jtype, hmp in enumerate(hmap_pred): suppressed_hmap = nms3d_cuda.NMSFilter3d( torch.nn.ConstantPad3d(1, 0)(hmp), 3, 1) nonzero_coords = torch.nonzero(suppressed_hmap).cpu() for coord in nonzero_coords: confidence = suppressed_hmap[tuple(coord)] if confidence > cnf.nms_th: pseudo3d_coords_pred.append( (jtype, coord[0].item(), coord[1].item(), coord[2].item())) confidences.append(confidence.cpu()) # pseudo-3D coordinates --> [reverse projection] --> real 3D coordinates coords3d_pred = [] for i in range(len(pseudo3d_coords_pred)): joint_type, cam_dist, y2d, x2d = pseudo3d_coords_pred[i] x2d, y2d, cam_dist = utils.rescale_to_real(x2d, y2d, cam_dist, q=cnf.q) x3d, y3d, z3d = utils.to3d(x2d, y2d, cam_dist, fx=fx, fy=fy, cx=cx, cy=cy) coords3d_pred.append((joint_type, x3d, y3d, z3d)) filter_joints(coords3d_pred, duplicate_th=0.05) # real 3D coordinates --> [association] --> list of poses poses = joint_association(coords3d_pred) # 3D poses -> [refiner] -> refined 3D poses refined_poses = [] for _pose in poses: refined_pose = refine_pose(pose=_pose, refiner=refiner) if refined_pose is not None: refined_poses.append(refined_pose) # show output print(f'\n\t▶▶ Showing results of \'{frame_path[0]}\'') print(f'\t▶▶ It may take some time: please wait') print(f'\t▶▶ Close mayavi window to continue') show_poses(refined_poses)
class CLeaR: def __init__(self): model = load_model() self.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") self.autoencoder = Autoencoder(model, self.device) self.predictor = Predictor(model, self.device) self.buffer_autoencoder_f = [] self.buffer_autoencoder_n = [] self.lim_autoencoder = 4000 self.threshold_autoencoder = 1 self.buffer_predictor_f = [] self.buffer_predictor_n = [] self.lim_predictor = 4000 self.threshold_predictor = 1 self.alpha = 0.95 def train_on_task(self, dataloader, warmup=False): for x, y in tqdm(dataloader, desc="Reading data stream", leave=None): x = x.to(self.device) y = [yi.to(self.device) for yi in y] with torch.no_grad(): z, x_ = self.autoencoder(x, True) error_autoencoder = mse(x, x_) if error_autoencoder > self.threshold_autoencoder or warmup: self.buffer_autoencoder_n.append(x) else: self.buffer_autoencoder_f.append(x) with torch.no_grad(): y_pred = self.predictor(z) error_predictor = minMSELoss(y_pred, y) if error_predictor > self.threshold_predictor or warmup: self.buffer_predictor_n.append((z, y)) else: self.buffer_predictor_f.append((z, y)) if len(self.buffer_autoencoder_n ) >= self.lim_autoencoder and not warmup: self.update_autoencoder() if len(self.buffer_predictor_n ) >= self.lim_predictor and not warmup: self.update_predictor() self.update_autoencoder() self.update_predictor() def update_autoencoder(self): self.autoencoder.update(self.buffer_autoencoder_n, self.buffer_autoencoder_f) with torch.no_grad(): x = torch.cat(self.buffer_autoencoder_n) x_ = self.autoencoder(x) error = mse(x, x_) self.threshold_autoencoder = self.alpha * error self.buffer_autoencoder_f = [] self.buffer_autoencoder_n = [] def update_predictor(self): self.predictor.update(self.buffer_predictor_n, self.buffer_predictor_f) with torch.no_grad(): z = torch.cat([zy[0] for zy in self.buffer_predictor_n]) y = [zy[1][0] for zy in self.buffer_predictor_n] y_pred = self.predictor(z) error = minMSELoss(y_pred, y) self.threshold_predictor = self.alpha * error self.buffer_predictor_f = [] self.buffer_predictor_n = [] def save_alexnet(self, path): d1 = self.autoencoder.encoder.state_dict() d1 = {'features.' + key: d1[key] for key in d1} d2 = self.predictor.model.state_dict() d = {**d1, **d2} torch.save(d, path) print("Saved model to {}".format(path))
def train_AE(args, dataset, model, loss_fn, config, num_overfit=-1, resume_optim=None, input_size=(1, 28, 28)): lr = config.getfloat('HYPERPARAMS', 'lr') print("Learning rate is %f" % lr) """ Create optimizer """ opt_list = list(model.parameters()) if args.model == 'VAE': #kl_loss_weight = torch.tensor(float(config['VAE_loss_weight'])) weighted_loss = AE.WeightedMultiLoss(init_values=[ config.getfloat('HYPERPARAMS', 'IMG_loss_weight'), config.getfloat('HYPERPARAMS', 'VAE_loss_weight') ], learn_weights=config.getboolean( 'HYPERPARAMS', 'learn_loss_weights')) opt_list.extend(list(weighted_loss.parameters())) auxillary = config.getboolean('HYPERPARAMS', 'auxillary', fallback=False) if auxillary and args.data != 'MNIST': raise RuntimeError('So far only MNIST has labels available') elif auxillary: n_classes = 10 if args.data == 'MNIST' else None aux_loss_fn = torch.nn.BCELoss() aux_loss_weight = AE.WeightedMultiLoss(init_values=[0, 0], learn_weights=False) opt_list.extend(list(aux_loss_weight.parameters())) loss_aux = [] if config.get('HYPERPARAMS', 'optimizer') == 'SGD': optim = torch.optim.SGD(params=opt_list, lr=lr) elif config.get('HYPERPARAMS', 'optimizer') == 'Adam': optim = torch.optim.Adam(params=opt_list, lr=lr, betas=(0.9, 0.99)) else: raise NotImplementedError('Unknown optimizer!') ##init model #AE.init(model) if resume_optim is not None: checkpoint = torch.load(resume_optim) optim.load_state_dict(checkpoint['optimizer_state_dict']) model.load_state_dict(checkpoint['model_state_dict']) print("Model/Optimizer loaded") """ Initialize training parameters """ overfit = num_overfit > -1 save_interval = config.getint('TRAINING', 'save_epoch_interval', fallback=10) ##create save folder timestamp = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M') path = args.model + '_' + args.data + '_' + timestamp if overfit: path += '_overfitted' save_dir = 'trained_models/' + path try: if not os.path.exists(save_dir): os.makedirs(save_dir) shutil.copyfile(args.config, save_dir + '/settings.ini') except OSError: print('Error: Creating directory. ' + save_dir) print("Created save directory at %s" % save_dir) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if torch.cuda.is_available(): print("CUDA available") if args.model == 'VAE': #kl_loss_weight = kl_loss_weight.cuda() weighted_loss.to(device) if auxillary: aux_loss_weight.to(device) if resume_optim is not None: for state in optim.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.to(device) model.to(device) """ Initialize visdom """ use_vis = config.getboolean('TRAINING', 'visdom', fallback=True) if use_vis: print('Use visdom') vis_env = 'training/{:s}_{:s}_{:s}'.format(args.model, args.data, timestamp) vis = visdom.Visdom() assert vis.check_connection( timeout_seconds=3), 'No connection could be formed quickly' plt_dict = dict(name='training loss', ytickmax=10, xlabel='epoch', ylabel='loss', legend=['train_loss']) vis_plot = vis.line(Y=[0], env=vis_env, opts=plt_dict) vis_image = vis.image(np.zeros(input_size), env=vis_env) vis_data = vis.image(np.zeros(input_size), env=vis_env) if args.model == 'VAE': vae_dict = dict(ytickmax=10, legend=['reconstruction loss', 'kl loss'], xlabel='epoch', ylabel='loss') vis_vae = vis.line(Y=[0, 0], env=vis_env, opts=vae_dict) vae_w_dict = dict(ymin=0, ymax=10, legend=['Weight Image', 'Weight KL-div'], xlabel='epoch', ylabel='value') vis_weights = vis.line(Y=[0, 0], env=vis_env, opts=vae_w_dict) if auxillary: aux_dict = dict(legend=['loss', 'aux loss'], xlabel='epoch', ylabel='loss') vis_aux = vis.line(Y=[0, 0], env=vis_env, opts=aux_dict) #vae_w_dict = dict(ymin=0, ymax=10, legend=['Weight Loss', 'Weight KL-div'], xlabel='epoch', ylabel='value') #vis_weights = vis.line(Y=[0,0], env=vis_env, opts=vae_w_dict) log_every = max( 1, int( len(dataset) / config.getint('TRAINING', 'log_every_dataset_chunk', fallback=5))) if not overfit else num_overfit if 'epochs' not in config['HYPERPARAMS']: warnings.warn( 'Number of epochs not specified in config - Use 20 as default', Warning) epochs = config.getint('HYPERPARAMS', 'epochs', fallback=20) loss_log = [] loss_vae = [] weight = [] x = [] div = 1. / float(len(dataset) if not overfit else num_overfit) t1 = None """ Actual training loop """ for epoch in tqdm.tqdm(range(epochs), total=epochs, desc='Epochs'): #print("Starting Epoch %d/%d (%s seconds)"%(epoch+1,epochs, 'N/A' if t1 is None else '{:.2f}'.format(time.time()-t1))) t1 = time.time() for i, (data, target) in tqdm.tqdm( enumerate(dataset), total=num_overfit if overfit else len(dataset), leave=False): if overfit: if i >= num_overfit: break data = data.to(device) loss = None optim.zero_grad() if args.model == 'VAE': prediction, mean, log_var, c = model(data) loss_img = loss_fn(prediction, data) ## KL = sum_i sigma_i^2 + mu_i^2 - log(sigma_i) -1 kl_div = 0.5 * torch.mean(mean**2 + torch.exp(log_var) - log_var - 1.0) #loss = loss_img + kl_loss_weight*kl_div#/(2*torch.exp(2*kl_loss_weight))+kl_loss_weight l = weighted_loss([loss_img, kl_div]) loss = sum(l) else: prediction, c = model(data) loss = loss_fn(prediction, data) if auxillary: target = torch.zeros( (data.size(0), n_classes), device=device).scatter_( 1, target.view(target.size(0), 1).to(device), 1) aux_loss = aux_loss_fn(c, target) l = aux_loss_weight([aux_loss, loss]) loss += sum(l) loss.backward() optim.step() if auxillary: loss_aux.append([loss.detach().cpu(), aux_loss.detach().cpu()]) if args.model == 'VAE': loss_vae.append( [loss_img.detach().cpu(), kl_div.detach().cpu()]) weight.append([ weighted_loss.weight0.detach().cpu(), weighted_loss.weight1.detach().cpu() ]) #print('Prediction\tmin:{:.2f}\tmax:{:.2f}\tmean:{:.2f}'.format(prediction.min(), prediction.max(), prediction.mean())) #print('data\tmin:{:.2f}\tmax:{:.2f}\tmean:{:.2f}'.format(data.min(), data.max(), data.mean())) loss_log.append(loss.detach().cpu()) x.append(epoch + float(i) * div) if i % log_every == 0: #print(" Loss after %d/%d iterations: %f"%(i,len(dataset) if not overfit else num_overfit,loss)) if use_vis: vis_plot = vis.line(win=vis_plot, X=x, Y=loss_log, env=vis_env, opts=plt_dict) img = prediction.cpu() img = img if img.shape[0] == 1 else img[0] data_img = torch.clamp(data.cpu(), 0, 1) data_img = data_img if data_img.shape[ 0] == 1 else data_img[0] vis_image = vis.image(img, win=vis_image, env=vis_env) vis_data = vis.image(data_img, win=vis_data, env=vis_env) if args.model == 'VAE': vis_vae = vis.line(win=vis_vae, X=x, Y=loss_vae, env=vis_env, opts=vae_dict) vis_weights = vis.line(win=vis_weights, X=x, Y=weight, env=vis_env, opts=vae_w_dict) if auxillary: vis_aux = vis.line(win=vis_aux, X=x, Y=loss_aux, env=vis_env, opts=aux_dict) vis.save(envs=[vis_env]) if ((epoch + 1) % save_interval) == 0: save_model(model, optim, save_dir, 'epoch_' + str(epoch + 1)) save_model(model, optim, save_dir, 'final_model') return optim
parser.add_argument('--embedding-dim', type=int, default=200) parser.add_argument('--enc-hidden-dim', type=int, default=100) parser.add_argument('--dec-hidden-dim', type=int, default=600) parser.add_argument('--interval', type=int, default=10) parser.add_argument('--cuda', type=bool, default=torch.cuda.is_available()) args = parser.parse_args() print(args) torch.manual_seed(args.seed) torch.backends.cudnn.deterministic = True train_loader, vocab = load(args.batch_size, args.seq_len) autoencoder = Autoencoder(args.enc_hidden_dim, args.dec_hidden_dim, args.embedding_dim, args.latent_dim, vocab.size(), args.dropout, args.seq_len) autoencoder.load_state_dict( torch.load('autoencoder.th', map_location=lambda x, y: x)) generator = Generator(args.n_layers, args.block_dim) critic = Critic(args.n_layers, args.block_dim) g_optimizer = optim.Adam(generator.parameters(), lr=args.lr) c_optimizer = optim.Adam(critic.parameters(), lr=args.lr) if args.cuda: autoencoder = autoencoder.cuda() generator = generator.cuda() critic = critic.cuda() print('G Parameters:', sum([p.numel() for p in generator.parameters() if \
AUTOENCODER = 'autoencoder_H18M05S37_04-23-20.pt' batch_size = 10 num_epochs = 20 epoch_size = 1000 n_frames = 30 n_vid_features = 3600 n_aud_features = 1 n_head = 3 n_layers = 3 dim_feedforward = 128 lr = 0.01 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') autoencoder = Autoencoder() autoencoder.load_state_dict(torch.load(AUTOENCODER)) autoencoder.to(device) autoencoder.eval() model = Classifier(n_vid_features, n_aud_features, n_head, n_layers, dim_feedforward) model = model.to(device) criterion = nn.BCEWithLogitsLoss() optimizer = optim.Adam(model.parameters(), lr=lr) start_time = datetime.datetime.now() print(f'start time: {str(start_time)}') print(f'using device: {device}') train_dataset = EncodedDeepfakeDataset(TRAIN_FOLDERS,
print('===> Using GPU to train') else: print('===> Using CPU to train') torch.manual_seed(args.seed) if args.cuda: torch.cuda.manual_seed(args.seed) print('===> Loaing datasets') images_A = get_image_paths("data/trump") images_B = get_image_paths("data/cage") images_A = load_images(images_A) / 255.0 images_B = load_images(images_B) / 255.0 images_A += images_B.mean(axis=(0, 1, 2)) - images_A.mean(axis=(0, 1, 2)) model = Autoencoder() print('===> Try resume from checkpoint') if os.path.isdir('checkpoint'): try: checkpoint = torch.load('./checkpoint/autoencoder.t7') model.load_state_dict(checkpoint['state']) start_epoch = checkpoint['epoch'] print('===> Load last checkpoint data') except FileNotFoundError: print('Can\'t found autoencoder.t7') else: start_epoch = 0 print('===> Start from scratch') if args.cuda:
sys.path.append("../model_builder") from Dataset_train import CustomTrainDataset import os import torch from skimage import io, transform import numpy as np from torch import nn import matplotlib.pyplot as plt from torch.utils.data import Dataset, DataLoader import matplotlib.pyplot as plt from models import Autoencoder from torchvision import transforms, utils import torch.nn.functional as F learning_rate = 0.0005 model = Autoencoder() criterion = nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) my_dataset = CustomTrainDataset(path='../data') dataloader = DataLoader(my_dataset, batch_size=4, shuffle=True) for epoch in range(40): for batch_idx, batch in enumerate(dataloader): # don't need labels, only the images (features) image = batch[0] out = model(image) model.encoder(image) cost = criterion(out, image) optimizer.zero_grad() cost.backward()
# compute model quality as `reward_map_fn`(validation accuracy) reward_map_fn = eval(reward_map_fn_str) torch.manual_seed(random_seed) np.random.seed(random_seed) device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device( 'cpu') dimension = 3 reps = 3 encoder_layers = [(1, 4), (4, 8), (8, 16), (16, 32), (32, 64)] # [(output_features, input_features), ...] decoder_layers = [(32, 128), (16, 64), (8, 32), (4, 16), (1, 8), (1, 2, False)] model = Autoencoder(dimension, reps, encoder_layers, decoder_layers, unet=True, use_sparsify=False) evaluator = model.encoder evaluator.freeze() evaluator.to(device) controller = FeatureExtractorController(encoder=evaluator, encoder_features=encoder_layers, N=4, latent_space_size=torch.LongTensor( [4, 4, 4]), latent_space_channels=64, archspace_epochs=10, reward_map_fn=reward_map_fn, device=device)
import torch import torch.nn as nn import numpy as np from models import Autoencoder, Generator from dataset import Corpus ##################### # Generating data ##################### ds = Corpus() vocab = ds.vocab generator = Generator(20, 100) generator.eval() generator.load_state_dict(torch.load('generator.th', map_location='cpu')) autoencoder = Autoencoder(100, 600, 200, 100, vocab.size(), 0.5, 22) autoencoder.eval() autoencoder.load_state_dict(torch.load('autoencoder.th', map_location='cpu')) # sample noise noise = torch.FloatTensor(np.random.normal(0, 1, (1, 100))) z = generator(noise[None,:,:]) # create new sent logits = autoencoder.decode(z).squeeze() seq = logits.argmax(dim=0) print(ds.decode(seq))
noise = "Gaussian" std_dev = 0.5 n_samples_vec = np.round(np.linspace(32, 16000, 100)).astype(int) for n_samples in n_samples_vec: print("Running with {} samples".format(n_samples)) session = keras.backend.get_session() dataset = images_dataset(batch_size=32, n_train_samples=n_samples, noise_type=noise, intensity=std_dev) ae_pol0 = Autoencoder(dataset, n_epochs=n_epochs, n_stages=20, test_runs=1, lr=lr, decay=decay, model_name="size:{}_pol0".format(n_samples)) ae_npol = Autoencoder(dataset, n_epochs=n_epochs, n_stages=20, test_runs=1, lr=lr, decay=decay, model_name="size:{}_npol".format(n_samples)) ae_pol0.train(do_valid=False, policy=policy) ae_npol.train(do_valid=False) def evaluate_model(dataset, model_pol, model_no_pol, batch_size=32):
cudnn.benchmark = True else: print('===> Using CPU to train') torch.manual_seed(args.seed) if args.cuda: torch.cuda.manual_seed(args.seed) print('===> Loaing datasets') images_A = get_image_paths("data/trump") images_B = get_image_paths("data/cage") images_A = load_images(images_A) / 255.0 images_B = load_images(images_B) / 255.0 images_A += images_B.mean(axis=(0, 1, 2)) - images_A.mean(axis=(0, 1, 2)) model = Autoencoder().to(device) print('===> Try resume from checkpoint') if os.path.isdir('checkpoint'): try: checkpoint = torch.load('./checkpoint/autoencoder.t7') model.load_state_dict(checkpoint['state']) start_epoch = checkpoint['epoch'] print('===> Load last checkpoint data') except FileNotFoundError: print('Can\'t found autoencoder.t7') else: start_epoch = 0 print('===> Start from scratch') criterion = nn.L1Loss()
policy = json2policy(policy_json) n_epochs = 500 batch_size = 32 n_train_samples = 50000 lr = 0.02 decay = 0.1 mnist_gauss_05 = MNIST(batch_size=32, n_train_samples=50, noise_type="Gaussian", intensity=0.5) ae_pol0 = Autoencoder(mnist_gauss_05, n_epochs=n_epochs, n_stages=20, lr=lr, decay=decay) ae_npol = Autoencoder(mnist_gauss_05, n_epochs=n_epochs, n_stages=20, lr=lr, decay=decay) ae_pol0.train(do_valid=False, policy=policy) ae_npol.train(do_valid=False, policy=policy) print("Performance with policy:") print(ae_pol0.test()) print("Performance without policy:") print(ae_npol.test())
return_energy=True) val = MuonPose("/home/jack/classes/thesis/autohas/LArCentroids/val/", return_energy=True) now = datetime.now().strftime("%b-%d-%y_%H:%M:%S") logger = SummaryWriter(log_dir=f"autoencoder_run/{now}/") device = 'cuda:0' dimension = 3 reps = 3 encoder_layers = [(1, 4), (4, 8), (8, 16), (16, 32), (32, 64)] # [(output_features, input_features), ...] decoder_layers = [(32, 128), (16, 64), (8, 32), (4, 16), (1, 8), (1, 2, False)] model = Autoencoder(dimension, reps, encoder_layers, decoder_layers, unet=True, use_sparsify=False) model.to_(device) loss_fn = torch.nn.MSELoss() optimizer = torch.optim.AdamW(model.parameters()) densifier = scn.Sequential().add( scn.InputLayer(3, torch.LongTensor([128, 128, 128]))).add(scn.SparseToDense(3, 1)) print("Want to test or load anything?") import code code.interact(local=locals())
parser.add_argument('--latent-dim', type=int, default=100) parser.add_argument('--enc-hidden-dim', type=int, default=100) parser.add_argument('--dec-hidden-dim', type=int, default=600) parser.add_argument('--interval', type=int, default=10) parser.add_argument('--cuda', type=bool, default=torch.cuda.is_available()) args = parser.parse_args() print(args) torch.manual_seed(args.seed) torch.backends.cudnn.deterministic = True train_loader, vocab = load(args.batch_size, args.seq_len) model = Autoencoder(args.enc_hidden_dim, args.dec_hidden_dim, args.embedding_dim, args.latent_dim, vocab.size(), args.dropout, args.seq_len) if args.cuda: model = model.cuda() print('Parameters:', sum([p.numel() for p in model.parameters() if \ p.requires_grad])) print('Vocab size:', vocab.size()) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=args.lr) best_loss = np.inf for epoch in range(1, args.epochs + 1):
def compute(exp_name): # type: (str) -> None cnf = Conf(exp_name=exp_name) # init Code Predictor predictor = CodePredictor() # type: BaseModel predictor.to(cnf.device) predictor.eval() predictor.requires_grad(False) predictor.load_w(cnf.exp_log_path / 'best.pth') # init Decoder autoencoder = Autoencoder() # type: BaseModel autoencoder.to(cnf.device) autoencoder.eval() autoencoder.requires_grad(False) autoencoder.load_w(Path(__file__).parent / 'models/weights/vha.pth') # init Hole Filler hole_filler = Refiner(pretrained=True) hole_filler.to(cnf.device) hole_filler.eval() hole_filler.requires_grad(False) hole_filler.load_w( Path(__file__).parent / 'models/weights/pose_refiner.pth') # init data loader ts = JTATestingSet(cnf=cnf) loader = DataLoader(dataset=ts, batch_size=1, shuffle=False, num_workers=0) metrics_dict = {} for th in THS: for key in ['pr', 're', 'f1']: metrics_dict[f'{key}@{th}'] = [] # without refinement metrics_dict[f'{key}@{th}+'] = [] # with refinement for step, sample in enumerate(loader): x, coords3d_true, fx, fy, cx, cy, frame_path = sample x = x.to(cnf.device) coords3d_true = json.loads(coords3d_true[0]) fx, fy, cx, cy = fx.item(), fy.item(), cx.item(), cy.item() # image --> [code_predictor] --> code code_pred = predictor.forward(x).unsqueeze(0) # code --> [decoder] --> hmap hmap_pred = autoencoder.decode(code_pred).squeeze() # hmap --> [local maxima search] --> pseudo-3D coordinates coords2d_pred = [] confs = [] for jtype, hmp in enumerate(hmap_pred.squeeze()): res = nms3d_cuda.NMSFilter3d(nn.ConstantPad3d(1, 0)(hmp), 3, 1) nz = torch.nonzero(res).cpu() for el in nz: confid = res[tuple(el)] if confid > 0.1: coords2d_pred.append( (jtype, el[0].item(), el[1].item(), el[2].item())) confs.append(confid.cpu()) # pseudo-3D coordinates --> [to_3d] --> real 3D coordinates coords3d_pred = [] for i in range(len(coords2d_pred)): joint_type, cam_dist, y2d, x2d = coords2d_pred[i] x2d, y2d, cam_dist = utils.rescale_to_real(x2d, y2d, cam_dist, q=cnf.q) x3d, y3d, z3d = utils.to3d(x2d, y2d, cam_dist, fx=fx, fy=fy, cx=cx, cy=cy) coords3d_pred.append((joint_type, x3d, y3d, z3d)) # real 3D coordinates --> [association] --> list of poses poses = coords_to_poses(coords3d_pred, confs) # a solitary joint is a joint that has been excluded from the association # process since no valid connection could be found; # note that only solitary joints with a confidence value >0.6 are considered all_pose_joints = [] for pose in poses: all_pose_joints += [(j.type, j.confidence, j.x3d, j.y3d, j.z3d) for j in pose] coords3d_pred_ = [(c[0], confs[k], c[1], c[2], c[3]) for k, c in enumerate(coords3d_pred)] solitary = [(s[0], s[2], s[3], s[4]) for s in (set(coords3d_pred_) - set(all_pose_joints)) if s[1] > 0.6] # list of poses --> [hole filler] --> refined list of poses refined_poses = [] for person_id, pose in enumerate(poses): confidences = [j.confidence for j in pose] pose = [(joint.type, joint.x3d, joint.y3d, joint.z3d) for joint in pose] refined_pose = hole_filler.refine(pose=pose, hole_th=0.2, confidences=confidences, replace_th=1) refined_poses.append(refined_pose) # refined list of poses --> [something] --> refined_coords3d_pred refined_coords3d_pred = [] for pose in refined_poses: refined_coords3d_pred += pose # compute metrics without refinement for th in THS: __m = joint_det_metrics(points_pred=coords3d_pred, points_true=coords3d_true, th=th) for key in ['pr', 're', 'f1']: metrics_dict[f'{key}@{th}'].append(__m[key]) # compute metrics with refinement for th in THS: __m = joint_det_metrics(points_pred=refined_coords3d_pred + solitary, points_true=coords3d_true, th=th) for key in ['pr', 're', 'f1']: metrics_dict[f'{key}@{th}+'].append(__m[key]) # print test progress print(f'\r>> processing test image {step} of {len(loader)}', end='') print('\r', end='') for th in THS: print(f'(PR, RE, F1)@{th}:' f'\tno_ref=(' f'{np.mean(metrics_dict[f"pr@{th}"]) * 100:.2f}, ' f'{np.mean(metrics_dict[f"re@{th}"]) * 100:.2f}, ' f'{np.mean(metrics_dict[f"f1@{th}"]) * 100:.2f})' f'\twith_ref=(' f'{np.mean(metrics_dict[f"pr@{th}+"]) * 100:.2f}, ' f'{np.mean(metrics_dict[f"re@{th}+"]) * 100:.2f}, ' f'{np.mean(metrics_dict[f"f1@{th}+"]) * 100:.2f}) ')
datetime.now().strftime('%Y-%m-%d_%H:%M:%S') ) dataset = getattr( __import__('datasets'), args.dataset.capitalize() + 'Dataset' ) train_dataset = dataset('train', args) val_dataset = dataset('validation', args) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True ) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False ) autoencoder = Autoencoder(args.encoder_layers, args.decoder_layers) classifier = LogisticRegression(args.encoder_layers[-1]) oracle = DL2_Oracle( learning_rate=args.dl2_lr, net=autoencoder, use_cuda=torch.cuda.is_available(), constraint=ConstraintBuilder.build( autoencoder, train_dataset, args.constraint ) ) binary_cross_entropy = nn.BCEWithLogitsLoss( pos_weight=train_dataset.pos_weight('train') if args.balanced else None ) optimizer = torch.optim.Adam( list(autoencoder.parameters()) + list(classifier.parameters()),
def main(): # Parse arguments args = get_parser().parse_args() # Check for GPU is_gpu = len(tf.config.list_physical_devices("GPU")) > 0 if not is_gpu and not args.allow_cpu: raise ValueError( "Cannot run the code on CPU. Please enable GPU support or pass the '--allow-cpu' flag." ) # Check if dataset should be recreated if not os.path.exists("cleaned_dataset") or args.force_recreate: prepare_data() # ........................................................................ # Let's select the train folds and the holdout fold # ........................................................................ # Grab folds folders (fold_1, ..., fold_k) folds_folders = glob.glob("patched_dataset/**") # Randomly select 1 fold for being the holdout final test holdout_fold = np.random.choice(folds_folders, 1)[0] # Keep the rest for training and performing k-fold # Search for elements of `folds_folders` that are not in `holdout_fold` train_folds = np.setdiff1d(folds_folders, holdout_fold) print(f"Train folds: {train_folds}") print(f"Holdout fold: {holdout_fold}") for k, fold in enumerate(train_folds): # Print current fold print(f"Train Fold {k+1}:") print_fold_stats(fold) # Now for the holdout fold print("Holdout Fold:") print_fold_stats(holdout_fold) # Generate the bands to keep bands_to_keep = np.round(np.linspace(0, 272 - 1, args.bands)).astype(int) # Load test images and labels print("Loading holdout fold images") images, labels = load_fold(holdout_fold, bands_to_keep) test_data = (images, labels) # Creat list for holding the training datasets folds train_data_list = [] for k, fold in enumerate(train_folds): print(f"Loading images of training fold {k+1}") # Load the normalized images to list images, labels = load_fold(fold, bands_to_keep) # Save images in dict with labels as value train_data_list.append([images, labels]) # Check if a baseline model should be created if args.baseline: # ........................................................................ # Let's now establish a baseline model # ........................................................................ model = BaselineModel().cross_validate( train_data_list, fully_connected_size=128, add_dropout_layers=True, ) # How does the model perform on unseen data? result = model.evaluate(test_data[0], test_data[1], batch_size=64) result = dict(zip(model.metrics_names, result)) print(result) # ........................................................................ # Let's now try to use the Autoencoder # ........................................................................ # Get train data (We only care about the features now) x_train, _ = build_train_data(train_data_list, range(len(train_data_list))) x_test = test_data[0] print("Training an Autoencoder") # Instantiate autoencoder autoencoder = Autoencoder(x_train[0].shape, n_bands=args.bands) # Prepare callbacks # 1. Reduce learning rate when loss is on plateau reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor="val_loss", factor=0.2, patience=15, min_lr=9e-4, verbose=1) # 2. Early stopping early_stop = tf.keras.callbacks.EarlyStopping( monitor="val_loss", min_delta=0, patience=20, mode="auto", verbose=0, restore_best_weights=True, ) # Compile the autoencoder autoencoder.compile( optimizer=tf.keras.optimizers.Adam(lr=1e-3), loss="mse", ) # Train the model history = {} history["autoencoder"] = autoencoder.fit( x_train, x_train, validation_split=0.2, batch_size=16, epochs=250, verbose=1, callbacks=[reduce_lr, early_stop], ) # Plot the Autoencoder loss curve plotter = tfdocs.plots.HistoryPlotter(metric="loss", smoothing_std=10) plotter.plot(history)