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)