def plot_centroid(args, model, data_loader, mus, filename=None, n_clusters=10):
if filename is None:
kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(mus)
filename = args.dataset + "-kmeans-" + str(args.latent_dim) + "-beta-" + str(args.beta) + ".pt"
path = os.path.join(args.save_dir, filename)
print("Saving kmeans: %s\n" % (path))
pickle.dump(kmeans, open(path, "wb"))
else:
#kmeans = KMeans(n_clusters=n_clusters, random_state=0)
path = os.path.join(args.save_dir, filename)
print("Loading kmeans: %s\n" % (path))
kmeans = pickle.load(open(path, "rb"))
centers = kmeans.cluster_centers_
model.eval()
device = get_device()
with torch.no_grad():
ytrue = torch.tensor(np.array([]), dtype=torch.long)
ypred = torch.tensor(np.array([]), dtype=torch.long)
ytrue = ytrue.to(device)
ypred = ypred.to(device)
for i, (x, target) in enumerate(data_loader):
x = x.to(device)
# print('x shape', x.shape)
target = target.to(device)
mu, logvar = model.encoder(x)
labels = kmeans.predict(mu.cpu())
#print('labels shape', labels.shape)
#y = to_categorical(torch.from_numpy(labels).to(device, dtype=torch.long), n_clusters=n_clusters)
y = torch.from_numpy(labels).to(device, dtype=torch.long)
ytrue = torch.cat((ytrue, target), axis=0)
ypred = torch.cat((ypred, y.to(device)), axis=0)
accuracy = unsupervised_labels(ytrue.cpu(), ypred.cpu(), n_clusters, n_clusters)
method = "Unsupervised (VAE)"
progress_bar(i,
len(data_loader),
'%s | Acc: %0.2f%%'
% (method, accuracy))
if i==0:
centroid_mu = centers[labels]
centroid_mu = torch.tensor(centroid_mu).to(device)
# print('mu and logvar shapes', centroid_mu.shape, logvar.shape)
z = model.reparameterize(centroid_mu, logvar)
centroid = model.decoder(z)
n = min(x.size(0), 64)
folder = "results"
os.makedirs(folder, exist_ok=True)
xp, mu, logvar = model(x)
img = torch.cat([x[:n],
xp.view(args.batch_size, 1, args.crop_size, args.crop_size)[:n]])
filename = folder + '/recon.png'
save_image(img.cpu(), filename, nrow=n)
img = torch.cat([x[:n],
centroid.view(args.batch_size, 1, args.crop_size, args.crop_size)[:n]])
filename = folder + '/centroid.png'
save_image(img.cpu(), filename, nrow=n)
def test(args, model, test_loader, epoch=0):
model.eval()
recon_losses = AverageMeter()
kl_losses = AverageMeter()
device = get_device()
with torch.no_grad():
for i, (x, _) in enumerate(test_loader):
x = x.to(device)
mu, logvar = model.encoder(x)
noise = torch.randn_like(mu) * 2 * logvar.exp().pow(0.5)
noised_mu = mu + noise
z = model.reparameterize(noised_mu, logvar)
noised_pos = model.decoder(z)
noise = torch.randn_like(mu) * 2 * logvar.exp().pow(0.5)
noised_mu = mu + noise
z = model.reparameterize(noised_mu, logvar)
noised_neg = model.decoder(z)
xp, mu, logvar = model(x)
loss, recon, kl = loss_function(args, xp, x, mu, logvar, args.beta)
recon_losses.update(recon.float().item() / x.size(0))
kl_losses.update(kl.float().item() / x.size(0))
progress_bar(
i, len(test_loader),
'Epoch: %d | Recon Loss: %.6f | KL Loss: %.6f' %
(epoch, recon_losses.avg, kl_losses.avg))
if i == 0:
n = min(x.size(0), 64)
folder = "results"
os.makedirs(folder, exist_ok=True)
img = torch.cat([
x[:n],
xp.view(args.batch_size, 1, args.crop_size,
args.crop_size)[:n]
])
filename = folder + '/recon.png'
save_image(img.cpu(), filename, nrow=n)
img = torch.cat([
x[:n],
noised_pos.view(args.batch_size, 1, args.crop_size,
args.crop_size)[:n]
])
filename = folder + '/noised_pos.png'
save_image(img.cpu(), filename, nrow=n)
img = torch.cat([
x[:n],
noised_neg.view(args.batch_size, 1, args.crop_size,
args.crop_size)[:n]
])
filename = folder + '/noised_neg.png'
save_image(img.cpu(), filename, nrow=n)
def get_dataloader(args, dataset):
device = get_device()
cuda = "cuda" in str(device)
kwargs = {'num_workers': 1, 'pin_memory': True} if cuda else {}
data_loader = DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
**kwargs)
return data_loader
def __init__(self, cfg):
super().__init__(cfg)
print_cuda_statistics()
self.device = get_device()
# define models
self.model
# define data_loader 1 or 2
# 1
# tr_dataset = custom_dataset(cfg.tr_data_pth)
# te_dataset = custom_dataset(cfg.te_data_pth)
# 2
# dataset = custom_dataset(cfg.data_pth)
tr_dataset, te_dataset = random_split(dataset, [train_size, test_size])
self.tr_loader = DataLoader(tr_dataset,
batch_size=cfg.bs,
shuffle=cfg.data_shuffle,
num_workers=cfg.num_w)
self.te_loader = DataLoader(te_dataset,
batch_size=cfg.bs,
shuffle=cfg.data_shuffle,
num_workers=cfg.num_w)
# define criterion
self.criterion = Loss()
# define optimizers for both generator and discriminator
self.optimizer = None
# initialize counter
self.current_epoch = 0
self.current_iteration = 0
self.best_metric = 0 # loss or accuracy or etc
# set the manual seed for torch
self.manual_seed = self.cfg.seed
if self.cuda:
torch.cuda.manual_seed_all(self.manual_seed)
torch.cuda.set_device(self.cfg.gpu_device)
self.model = self.model.cuda()
self.loss = self.loss.cuda()
self.logger.info("Program will run on *****GPU-CUDA***** ")
print_cuda_statistics()
else:
self.logger.info("Program will run on *****CPU*****\n")
# Model Loading from cfg if not found start from scratch.
self.load_checkpoint(self.cfg.checkpoint_file)
# Summary Writer
self.summary_writer = None
def plot_tsne2(args, mus, model, filename=None, n_clusters=10):
kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(mus)
x = mus[:,0]
y = mus[:,1]
data = pd.DataFrame()
data['x'] = x
data['y'] = y
data['label'] = kmeans.labels_
#current_palette = sns.color_palette()
#sns.palplot(current_palette)
# ax = sns.scatterplot(
# x="x", y="y",
# hue="label",
# data=data,
# palette=sns.color_palette("hls", n_clusters),
# alpha=0.3
# )
# fig = ax.get_figure()
# fig.savefig('results/tsne.png')
centers = kmeans.cluster_centers_
print('centers shape',centers[0].shape)
device = get_device()
centroid_idx = list(range(n_clusters))
# for i in range(n_clusters):
centroid_mu = centers[centroid_idx]
print('centroid shape',centroid_mu.shape)
centroid_mu = torch.tensor(centroid_mu).to(device)
# logvar = torch.zeros(centroid_mu.shape).to(device)
# print('mu and logvar shapes', centroid_mu.shape, logvar.shape)
# z = model.reparameterize(centroid_mu, logvar)
z = centroid_mu
centroid = model.decoder(z)
# print('i',i)
print('centroid', centroid.shape)
# cv2.imshow('image',centroid)
# n = min(x.size(0), 64)
folder = "results"
# os.makedirs(folder, exist_ok=True)
# xp, mu, logvar = model(x)
# img = torch.cat([x[:n],
# xp.view(args.batch_size, 1, args.crop_size, args.crop_size)[:n]])
# filename = folder + '/centroid_images'+str(i)+'.png'
filename = folder + '/centroid_images.png'
save_image(centroid, filename,nrow=n_clusters)
def train(args, model, train_loader, optimizer, epoch):
model.train()
recon_losses = AverageMeter()
kl_losses = AverageMeter()
device = get_device()
for i, (x, _) in enumerate(train_loader):
x = x.to(device)
optimizer.zero_grad()
xp, mu, logvar = model(x)
loss, recon, kl = loss_function(args, xp, x, mu, logvar, args.beta)
loss.backward()
recon_losses.update(recon.float().item()/x.size(0))
kl_losses.update(kl.float().item()/x.size(0))
optimizer.step()
progress_bar(i,
len(train_loader),
'Epoch: %d | Recon Loss: %.6f | KL Loss: %.6f' % (epoch, recon_losses.avg, kl_losses.avg))
def __init__(self, config):
super().__init__(config)
self.config = config
self.device = get_device()
set_seeds(self.config.seed)
self.current_epoch = 1
self.global_step = 0
self.best_valid_mean_iou = 0
self.model = BERTModel4Pretrain(self.config)
self.criterion1 = nn.CrossEntropyLoss(reduction='none')
self.criterion2 = nn.CrossEntropyLoss()
self.optimizer = optim4GPU(self.config, self.model)
self.writer = SummaryWriter(log_dir=self.config.log_dir)
tokenizer = FullTokenizer(self.config, do_lower_case=True)
train_dataset = SentencePairDataset(self.config, tokenizer, 'train')
validate_dataset = SentencePairDataset(self.config, tokenizer,
'validate')
a = train_dataset.__getitem__(0)
self.train_dataloader = DataLoader(
train_dataset,
batch_size=self.config.batch_size,
num_workers=self.config.data_loader_workers,
pin_memory=self.config.pin_memory)
self.validate_dataloader = DataLoader(
validate_dataset,
batch_size=self.config.batch_size,
num_workers=self.config.data_loader_workers,
pin_memory=self.config.pin_memory)
self.model = self.model.to(self.device)
if self.config.data_parallel:
self.model = nn.DataParallel(self.model)
self.load_checkpoint(self.config.checkpoint_to_load)
def tsne(args, model, data_loader, tsne=False):
model.eval()
device = get_device()
mus = np.array([])
with torch.no_grad():
for i, (x, _) in enumerate(data_loader):
x = x.to(device)
mu, logvar = model.encoder(x)
if len(mus) == 0:
mus = mu.cpu()
else:
mus = np.concatenate((mus, mu.cpu()), axis=0)
progress_bar(i, len(data_loader))
if tsne:
mus_tsne = TSNE(n_components=2).fit_transform(mus)
print(mus_tsne.shape)
return mus, mus_tsne
print(mus.shape)
return mus, _
def vae():
args = get_args()
device = get_device()
#torch.manual_seed(args.seed)
model = VAE(args.latent_dim, args.crop_size).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-3)
if args.dataset == "mnist":
dataset = datasets.MNIST('./data',
train=True,
download=True,
transform=crop_transform(args.crop_size))
train_loader = get_dataloader(args, dataset)
dataset = datasets.MNIST('./data',
train=False,
download=True,
transform=crop_transform(args.crop_size))
test_loader = get_dataloader(args, dataset)
elif args.dataset == "fashionmnist":
dataset = datasets.FashionMNIST('./data',
train=True,
download=True,
transform=crop_transform(
args.crop_size))
train_loader = get_dataloader(args, dataset)
dataset = datasets.FashionMNIST('./data',
train=False,
download=True,
transform=crop_transform(
args.crop_size))
test_loader = get_dataloader(args, dataset)
else:
exit(0)
print("Dataset: ", args.dataset)
folder = args.save_dir
os.makedirs(folder, exist_ok=True)
compute_kmeans = False
if args.train:
for epoch in range(1, args.epochs + 1):
print("Epoch %d" % (epoch))
train(args, model, train_loader, optimizer, epoch)
test(args, model, test_loader, epoch)
filename = args.dataset + "-vae-" + str(
args.latent_dim) + "-beta-" + str(args.beta) + ".pt"
path = os.path.join(folder, filename)
print("Saving weights: %s\n" % (path))
torch.save(model.state_dict(), path)
compute_kmeans = True
if args.restore_weights is not None:
path = os.path.join(folder, args.restore_weights)
print("Loading weights... ", path)
model.load_state_dict(torch.load(path))
test(args, model, test_loader)
compute_kmeans = True
if compute_kmeans:
mus, mus_tsne = tsne(args, model, train_loader, tsne=args.tsne)
if args.tsne:
plot_tsne(mus_tsne)
if args.kmeans is None:
plot_centroid(args, model, train_loader, mus, args.kmeans)
else:
plot_centroid(args, model, test_loader, mus, args.kmeans)
def __init__(self, cfg):
super().__init__(cfg)
self.cfg = cfg
self.device = get_device()
# define models
# inception_v3 input size = ( N x 3 x 299 x 299 )
self.model = inception_v3(pretrained=True, num_classes=cfg.num_classes)
# define data_loader
self.dataset = inception_data(cfg).get_dataset()
tr_size = int(cfg.train_test_ratio * len(self.dataset))
te_size = len(self.dataset) - tr_size
tr_dataset, te_dataset = random_split(self.dataset, [tr_size, te_size])
self.tr_loader = DataLoader(tr_dataset,
batch_size=cfg.bs,
shuffle=cfg.data_shuffle,
num_workers=cfg.num_workers)
self.te_loader = DataLoader(te_dataset,
batch_size=cfg.bs,
shuffle=cfg.data_shuffle,
num_workers=cfg.num_workers)
# define loss
self.loss = torch.tensor(0)
self.criterion = CrossEntropyLoss()
# define optimizers for both generator and discriminator
self.optimizer = RMSprop(self.model.parameters(), lr=cfg.lr)
# initialize counter
self.current_epoch = 0
self.current_iteration = 0
self.best_metric = 0
self.best_info = ""
# set cuda flag
self.is_cuda = torch.cuda.is_available()
if self.is_cuda and not self.cfg.cuda:
self.logger.info(
"WARNING: You have a CUDA device, so you should probably enable CUDA"
)
self.cuda = self.is_cuda & self.cfg.cuda
# set the manual seed for torch
self.manual_seed = self.cfg.seed
if self.cuda:
torch.cuda.manual_seed(self.manual_seed)
self.model = self.model.to(self.device)
if self.cfg.data_parallel:
self.model = nn.DataParallel(self.model)
self.logger.info("Program will run on *****GPU-CUDA***** ")
else:
self.model = self.model.to(self.device)
torch.manual_seed(self.manual_seed)
self.logger.info("Program will run on *****CPU*****\n")
# Model Loading from cfg if not found start from scratch.
self.exp_dir = os.path.join('./experiments', cfg.exp_name)
self.load_checkpoint(self.cfg.checkpoint_filename)
# Summary Writer
self.summary_writer = SummaryWriter(
log_dir=os.path.join(self.exp_dir, 'summaries'))
def __init__(self, cfg):
super().__init__(cfg)
self.cfg = cfg
self.device = get_device()
# define models
self.model = EAST(pretrained=True, cfg=cfg)
# set cuda flag
self.cuda = (self.device == torch.device('cuda')) and cfg.cuda
# set the manual seed for torch
self.manual_seed = cfg.seed
if self.cuda:
torch.cuda.manual_seed_all(self.manual_seed)
self.model = self.model.to(self.device)
if cfg.data_parallel:
self.model = nn.DataParallel(self.model)
self.logger.info("Program will run on *****GPU-CUDA***** ")
else:
self.model = self.model.to(torch.device('cpu'))
self.logger.info("Program will run on *****CPU*****\n")
# ready done for prediction
if cfg.mode == "predict":
self.load_checkpoint(cfg.checkpoint_file)
self.model.eval()
return
# define data_loader
train_dataset = custom_dataset(cfg.tr_im_pth, cfg.tr_gt_pth)
test_dataset = custom_dataset(cfg.te_im_pth, cfg.te_gt_pth)
self.train_loader = DataLoader(train_dataset,
batch_size=cfg.bs,
shuffle=cfg.data_shuffle,
num_workers=cfg.num_w,
drop_last=True)
self.test_loader = DataLoader(test_dataset,
batch_size=cfg.bs,
shuffle=cfg.data_shuffle,
num_workers=cfg.num_w,
drop_last=True)
# define loss
self.criterion = EASTLoss()
# define optimizers for both generator and discriminator
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=cfg.lr)
self.scheduler = lr_scheduler.MultiStepLR(self.optimizer,
milestones=[cfg.epochs // 2],
gamma=0.1)
# initialize counter
self.current_epoch = 1
self.current_iteration = 1
self.best_metric = np.inf # loss
self.best_info = ''
# Model Loading from the latest checkpoint if not found start from scratch.
self.exp_dir = os.path.join('./experiments', cfg.exp_name)
self.load_checkpoint(cfg.checkpoint_file)
# Summary Writer
self.summary_writer = SummaryWriter(
log_dir=os.path.join(self.exp_dir, 'summaries'))
def __init__(self, cfg):
super().__init__(cfg)
self.device = get_device()
cfg.device = self.device
self.cfg = cfg
# dataset
train_dataset = Im2LatexDataset(cfg, mode="train")
self.id2token = train_dataset.id2token
self.token2id = train_dataset.token2id
collate = custom_collate(self.token2id, cfg.max_len)
self.train_loader = DataLoader(train_dataset,
batch_size=cfg.bs,
shuffle=cfg.data_shuffle,
num_workers=cfg.num_w,
collate_fn=collate,
drop_last=True)
if cfg.valid_img_path != "":
valid_dataset = Im2LatexDataset(cfg,
mode="valid",
vocab={
'id2token': self.id2token,
'token2id': self.token2id
})
self.valid_loader = DataLoader(valid_dataset,
batch_size=cfg.bs //
cfg.beam_search_k,
shuffle=cfg.data_shuffle,
num_workers=cfg.num_w,
collate_fn=collate,
drop_last=True)
# define models
self.model = Im2LatexModel(cfg) # fill the parameters
# weight initialization setting
for name, param in self.model.named_parameters():
if 'localization_fc2' in name:
print(f'Skip {name} as it is already initialized')
continue
try:
if 'bias' in name:
torch.nn.init.constant_(param, 0.0)
elif 'weight' in name:
torch.nn.init.kaiming_normal_(param)
except Exception as e: # for batchnorm.
if 'weight' in name:
param.data.fill_(1)
continue
self.model = DataParallel(self.model)
# define criterion
self.criterion = cal_loss
self.optimizer = torch.optim.Adam(params=self.model.parameters(),
lr=cfg.lr,
betas=(cfg.adam_beta_1,
cfg.adam_beta_2))
milestones = cfg.milestones
self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer,
milestones,
gamma=cfg.gamma,
verbose=True)
# initialize counter
self.current_epoch = 1
self.current_iteration = 1
self.best_metric = 100
self.best_info = ''
# set the manual seed for torch
torch.cuda.manual_seed_all(self.cfg.seed)
if self.cfg.cuda:
self.model = self.model.to(self.device)
self.logger.info("Program will run on *****GPU-CUDA***** ")
else:
self.logger.info("Program will run on *****CPU*****\n")
# Model Loading from cfg if not found start from scratch.
self.exp_dir = os.path.join('./experiments', cfg.exp_name)
self.load_checkpoint(cfg.checkpoint_filename)
# Summary Writer
self.summary_writer = SummaryWriter(
log_dir=os.path.join(self.exp_dir, 'summaries'))
