def fix_seed(seed: int,
deterministic: bool = False,
benchmark: bool = False) -> None:
"""
Fixed the ``Seed`` value of PyTorch, NumPy, Pure Python Random at once.
Examples:
>>> import torch
>>> import numpy as np
>>> fix_seed(0)
>>> x = torch.randn(...)
>>> y = np.random.randn(...)
Args:
seed (int): random state (sedd)
deterministic (bool): Whether to ensure reproducibility as much as possible on CuDNN.
benchmark (bool):
Returns:
None
"""
std_seed(seed)
os_environ["PYTHONHASHSEED"] = str(seed)
np_seed(seed)
torch.manual_seed(seed)
if cuda_is_available():
if deterministic:
cudnn.deterministic = True
if benchmark:
cudnn.benchmark = False
def generate_images(model,
batch,
mask_descriptors,
num_samples=64,
temp=1.,
verbose=False):
"""Generates image completions based on the images in batch masked by the
masks in mask_descriptors. This will generate
batch.size(0) * len(mask_descriptors) * num_samples completions, i.e.
num_samples completions for every image and mask combination.
Parameters
----------
model : pixconcnn.models.pixel_constrained.PixelConstrained instance
batch : torch.Tensor
mask_descriptors : list of mask_descriptor
See utils.masks.MaskGenerator for allowed mask_descriptors.
num_samples : int
Number of samples to generate for a given image-mask combination.
temp : float
Temperature for sampling.
verbose : bool
If True prints progress information while generating images
"""
device = torch_device("cuda" if cuda_is_available() else "cpu")
model.to(device)
outputs = []
for i in range(batch.size(0)):
outputs_per_img = []
for j in range(len(mask_descriptors)):
if verbose:
print("Generating samples for image {} using mask {}".format(
i, mask_descriptors[j]))
# Get image and mask combination
img = batch[i:i + 1]
mask_generator = MaskGenerator(model.prior_net.img_size,
mask_descriptors[j])
mask = mask_generator.get_masks(1)
# Create conditional pixels which will be used to sample completions
cond_pixels = get_repeated_conditional_pixels(
img, mask, model.prior_net.num_colors, num_samples)
cond_pixels = cond_pixels.to(device)
samples, log_probs = model.sample(cond_pixels,
return_likelihood=True,
temp=temp)
outputs_per_img.append({
"orig_img": img,
"cond_pixels": cond_pixels,
"mask": mask,
"samples": samples,
"log_probs": log_probs
})
outputs.append(outputs_per_img)
return outputs
def select_device(gpu_if_possible: bool):
"""
Select the device handling computations based on whether a GPU is
requested and whether it is available.
"""
# employing a GPU if requested and possible:
if gpu_if_possible and cuda_is_available():
return 'cuda:0'
return 'cpu'
def __init__(
self,
dataset: Dataset,
instance: Instance,
model: Model,
optimizer: Optimizer,
scheduler: Any,
epochs: int,
batch_size: int,
run_dir_path: Path,
eval_every: int,
limit_epochs_at: Optional[int],
train_eval_split: float,
metric_names: List[str],
selection_metric: str,
kept_checkpoints: int,
cuda_device: Optional[int],
) -> None:
self.dataset = dataset
self.instance = instance
self.optimizer = optimizer
self.scheduler = scheduler
self.epochs = epochs
self.batch_size = batch_size
self.run_dir_path = run_dir_path
self.eval_every = eval_every
self.train_eval_split = train_eval_split
self.limit_epochs_at = 10e1000 if limit_epochs_at is None else limit_epochs_at
self.selection_metric = selection_metric
self.kept_checkpoints = kept_checkpoints
self._checkpoints_stats: List[Tuple[float, str]] = []
if cuda_device is not None:
if not cuda_is_available():
raise RuntimeError("CUDA is not available on this system.")
self.use_cuda = True
self.device = torch_device("cuda:%d" % cuda_device)
else:
self.use_cuda = False
self.device = torch_device("cpu")
self.model = model.to(self.device)
self._checkpoints_dir = self.run_dir_path / "checkpoints"
self._writers: Dict[DataType, SummaryWriter] = {}
self._accumulated_metrics: Dict[DataType, Dict[
Metric,
List[float]]] = defaultdict(lambda: defaultdict(lambda: []))
self._metrics = [_metrics[metric_name] for metric_name in metric_names]
self._metric_names = {
metric: metric_name
for metric, metric_name in zip(self._metrics, metric_names)
}
self._epochs_size = len(str(epochs))
def __init__(
self,
dpr_fn: str,
tokenizer_fn: str,
tokenizer_max_len: int,
):
self.dpr = DPRReader.from_pretrained(dpr_fn)
self.tokenizer_max_len = tokenizer_max_len
self.tokenizer = DPRReaderTokenizer.from_pretrained(
tokenizer_fn, max_len=tokenizer_max_len)
device = 'cuda' if cuda_is_available() else 'cpu'
self.dpr.to(device)
self.device = device
def make_results_reproducible(model_is_convolutional: bool = False) -> None:
"""
Make the subsequent instructions produce purely deterministic outputs
by fixing all the relevant seeds:
"""
random_seed(0)
_ = numpy_seed(0)
_ = torch_manual_seed(0)
# since GPU computations may introduce additional stochasticity with
# their convolutional operation optimizations:
if model_is_convolutional:
if cuda_is_available():
# disabling benchmarking and choosing among convolution operation
# implementation alternatives, which is stochastic based due to
# noise and hardware:
cudnn.benchmark = False
# ansuring deterministic algorithms for colvolutional operations
# are employed:
cudnn.deterministic = True
if vc_num == 512:
extractor = models.vgg16(pretrained=True).features
elif vc_num == 256:
extractor = models.vgg16(pretrained=True).features[:12]
elif vc_num == 128:
extractor = models.vgg16(pretrained=True).features[:10]
elif nn_type[:6] == 'resnet' or nn_type == 'resnext' or nn_type == 'alexnet':
layer = 'last' # 'last','second'
extractor = resnet_feature_extractor(nn_type, layer)
elif nn_type == 'unet':
layer = 'pool5'
path_to_unet = os.path.join(Directories.CHECKPOINTS, unet_filename)
unet = UNet(pretrained=True)
device = device('cuda:0' if cuda_is_available() else 'cpu')
unet.load_state_dict(
load(path_to_unet, map_location=device)['model_state_dict'])
unet_ones = UNet(pretrained=False)
unet_ones = unet_ones.get_features()
if vc_num == 1024:
extractor = unet.get_features()[:24]
elif vc_num == 512:
extractor = unet.get_features()[:19]
elif vc_num == 256:
extractor = unet.get_features()[:15]
elif vc_num == 128:
extractor = unet.get_features()[:10]
else:
from .draft import RelEnc as RelationalEncoder, Encoder, RelationHead, RelationHead2, RelationHead3, RelationHeadLong, RelationHeadCone, aggregate
from .from_paper import MultiCIFAR10, train_transform
from .datasets import MRCData, MRCSampler
from torch.utils.data import DataLoader
from torch import save as save_model
from torch.cuda import is_available as cuda_is_available
device = "cuda" if cuda_is_available() else "cpu"
def relation_head(i):
return [
RelationHead, RelationHead2, RelationHead3, RelationHeadLong,
RelationHeadCone
][i]
def train_classifiers(cfg,
annot_dataloader,
count_celltypes,
count_classes,
track=False,
test_data=None):
device = torch.device(
'cuda' if cuda_is_available() and cfg.use_cuda else 'cpu')
celltype_clf = Classifier(inp_size=cfg.classifier_input_dim,
out_size=count_celltypes)
form_clf = Classifier(inp_size=cfg.classifier_input_dim,
out_size=count_classes)
if cfg.use_cuda and cuda_is_available():
celltype_clf = celltype_clf.cuda()
form_clf = form_clf.cuda()
celltype_clf_opt = torch.optim.Adam(celltype_clf.parameters(),
weight_decay=cfg.celltype_clf_wdecay,
lr=cfg.celltype_clf_lr)
form_clf_opt = torch.optim.Adam(form_clf.parameters(),
weight_decay=cfg.form_clf_wdecay,
lr=cfg.form_clf_lr)
celltype_criterion = nn.CrossEntropyLoss()
form_criterion = nn.CrossEntropyLoss()
print('\n')
metrics = {
'celltype_losses': [],
'form_losses': [],
'celltype_accuracy': [],
'form_accuracy': [],
'epoch': []
}
for epoch in range(cfg.classifier_epochs):
if cfg.verbose == 'all':
print(f'\rTraining classifier [{epoch+1}/{cfg.classifier_epochs}]',
end='')
celltype_clf.train()
form_clf.train()
metrics['epoch'].append(epoch + 1)
celltype_av_loss = 0.
form_av_loss = 0.
counter = 0
for exp_, form_, cell_type_ in annot_dataloader:
exp_ = exp_.to(device)
form_ = form_.argmax(-1).to(device)
cell_type_ = cell_type_.argmax(-1).to(device)
predicted_celltype_ = celltype_clf(exp_)
celltype_loss_on_batch = celltype_criterion(
predicted_celltype_, cell_type_)
celltype_av_loss += celltype_loss_on_batch.item()
celltype_clf_opt.zero_grad()
celltype_loss_on_batch.backward(retain_graph=True)
celltype_clf_opt.step()
predicted_form_ = form_clf(exp_)
form_loss_on_batch = form_criterion(predicted_form_, form_)
form_av_loss += form_loss_on_batch.item()
form_clf_opt.zero_grad()
form_loss_on_batch.backward(retain_graph=True)
form_clf_opt.step()
counter += 1
metrics['celltype_losses'].append(celltype_av_loss / counter)
metrics['form_losses'].append(form_av_loss / counter)
if track:
celltype_clf.eval()
form_clf.eval()
counter = 0
celltype_accuracy = 0.
form_accuracy = 0.
for exp_, form_, cell_type_ in test_data:
exp_ = exp_.to(device)
form_ = form_.argmax(-1)
cell_type_ = cell_type_.argmax(-1)
predicted_celltype_ = celltype_clf(
exp_).cpu().detach().numpy().argmax(-1)
cell_type_ = cell_type_.cpu().detach().numpy()
celltype_accuracy += (predicted_celltype_ == cell_type_).mean()
predicted_form_ = form_clf(exp_).cpu().detach().numpy().argmax(-1)
form_ = form_.cpu().detach().numpy()
form_accuracy += (predicted_form_ == form_).mean()
counter += 1
metrics['celltype_accuracy'].append(celltype_accuracy / counter)
metrics['form_accuracy'].append(form_accuracy / counter)
if track:
return celltype_clf, form_clf, metrics
return celltype_clf, form_clf
def test(cfg,
vae_model,
discrim,
annot_dataloader,
test_expression,
class_ohe_test,
celltype_test,
pretrained_classifiers=None,
dataset_name: str = '') -> Dict[str, Union[float, str, None]]:
'''Calculate metrics and return dict
:Param cfg: Config dataclass
:Param vae_model: autoecoder with declarated signature
:Param discrim: deprecated, unused
:Param annot_dataloader: dataloader(expression, batch_indices, labels)
for classifier training
:Param test_expression:
'''
#STYLE TRANSFER
ge_transfer_raw = np.repeat(test_expression,
class_ohe_test.shape[1],
axis=0)
ge_transfer_raw = Tensor(ge_transfer_raw)
init_classes_transfer = np.repeat(class_ohe_test,
class_ohe_test.shape[1],
axis=0)
init_classes_transfer = Tensor(init_classes_transfer)
init_celltypes_transfer = np.repeat(celltype_test,
class_ohe_test.shape[1],
axis=0)
init_celltypes_transfer = Tensor(init_celltypes_transfer)
device = torch.device('cpu')
if cfg.use_cuda and cuda_is_available():
device = torch.device('cuda')
ge_transfer_raw = ge_transfer_raw.cuda()
init_classes_transfer = init_classes_transfer.cuda()
init_celltypes_transfer = init_celltypes_transfer.cuda()
target_classes_transfer = np.zeros(
(class_ohe_test.shape[0] * class_ohe_test.shape[1],
class_ohe_test.shape[1]))
target_classes_transfer[np.arange(target_classes_transfer.shape[0]),
np.arange(target_classes_transfer.shape[0]) %
target_classes_transfer.shape[1]] = 1
target_classes_transfer = Tensor(target_classes_transfer)
if cfg.use_cuda and cuda_is_available():
target_classes_transfer = target_classes_transfer.cuda()
transfer_expression_tensor = vae_model(ge_transfer_raw,
target_classes_transfer)[0]
if isinstance(transfer_expression_tensor, tuple):
transfer_expression_tensor = transfer_expression_tensor[0]
transfer_expression_np = transfer_expression_tensor.cpu().detach().numpy()
test_expression_tensor = Tensor(test_expression)
class_ohe_test_tensor = Tensor(class_ohe_test)
if cfg.use_cuda and cuda_is_available():
target_classes_transfer = target_classes_transfer.cuda()
test_expression_tensor = test_expression_tensor.cuda()
class_ohe_test_tensor = class_ohe_test_tensor.cuda()
reconstruction = vae_model(test_expression_tensor,
class_ohe_test_tensor)[0]
if isinstance(reconstruction, tuple):
reconstruction = reconstruction[0]
reconstruction = reconstruction.cpu().detach().numpy()
mse = (reconstruction - test_expression)**2
residual_transfer = ge_transfer_raw.cpu().numpy() - transfer_expression_np
res_norm = (residual_transfer**2).mean(1).reshape(
(-1, cfg.count_classes)) #l2 norm,
res_equal = class_ohe_test.argmax(1)
res_equal = res_equal if isinstance(
res_equal, np.ndarray) else res_equal.cpu().detach().numpy()
if cfg.verbose in ('all', 'test'):
print('Mean square error reconstructed expression:', np.mean(mse))
print('Mean values of test expression:', test_expression.mean())
print('Mean values of reconstructed expression:',
reconstruction.mean())
print('part of test expression <0.5 mean:',
(test_expression < 0.5).mean())
print('part reconstructed exrression <0.5 mean:',
(reconstruction < 0.5).mean())
print('Calibration accuracy', (res_norm.argmin(1) == res_equal).mean())
print('\n')
test_latents = vae_model.latents(test_expression_tensor,
class_ohe_test_tensor)
transfered_latents = vae_model.latents(ge_transfer_raw,
target_classes_transfer)
#train
celltype_clf_expr = None
celltype_clf_latents = None
form_clf_expr = None
form_clf_latents = None
if not isinstance(test_latents, torch.Tensor):
test_latents = torch.Tensor(test_latents)
if not isinstance(transfered_latents, torch.Tensor):
transfered_latents = torch.Tensor(transfered_latents)
if cfg.use_cuda and cuda_is_available():
test_latents = test_latents.cuda()
transfered_latents = transfered_latents.cuda()
if pretrained_classifiers is None:
cfg.classifier_input_dim = cfg.input_dim
celltype_clf_expr, form_clf_expr = train_classifiers(
cfg, annot_dataloader, celltype_test.shape[1], cfg.count_classes)
cfg.classifier_input_dim = vae_model.latent_dim
latents_dataset = torch.utils.data.TensorDataset(
test_latents, class_ohe_test_tensor,
torch.randn(test_latents.shape[0], cfg.count_labels))
latents_dataloader = torch.utils.data.DataLoader(
latents_dataset,
batch_size=cfg.batch_size,
shuffle=True,
drop_last=True)
celltype_clf_latents, form_clf_latents = train_classifiers(
cfg, latents_dataloader, celltype_test.shape[1], cfg.count_classes)
else:
celltype_clf = pretrained_classifiers[0]
form_clf = pretrained_classifiers[1]
if isinstance(test_latents, Tensor):
test_latents_np = test_latents.cpu().detach().numpy()
if isinstance(transfered_latents, Tensor):
transfered_latents_np = transfered_latents.cpu().detach().numpy()
# entropy batch mixing
ebm_score = {}
ebm_score['test'] = [0. for i in range(1)]
ebm_score['transfered'] = [0. for i in range(1)]
for n in range(1):
for lat, bind, key in zip(
(test_latents_np, transfered_latents_np),
(class_ohe_test_tensor, target_classes_transfer),
ebm_score.keys()):
batch_ind = bind.argmax(1).cpu().numpy()
ind = list(combinations(np.unique(batch_ind), 2))
for i in ind:
a = (batch_ind == i[0])
b = (batch_ind == i[1])
condition = np.logical_or(a, b)
#Important breakpoint
ebm_score['test'][n] += entropy_batch_mixing(
lat[condition], batch_ind[condition])
ebm_score['transfered'][n] += entropy_batch_mixing(
lat[condition], batch_ind[condition])
if len(ind) > 0:
ebm_score[key][n] /= len(ind)
# KNN Purity
test_purity = []
transfered_purity = []
test_purity.append(get_knn_purity(test_latents_np,
celltype_test.argmax(1)))
transfered_purity.append(
get_knn_purity(transfered_latents_np,
target_classes_transfer.argmax(dim=1).cpu().numpy()))
celltype_train_labelenc = None
if not annot_dataloader is None:
celltype_train_labelenc = annot_dataloader.dataset.tensors[2].argmax(
1).cpu().numpy()
celltype_test_raw_labelenc = celltype_test.argmax(1)
celltype_test_labelenc = init_celltypes_transfer.argmax(1).cpu().numpy()
#form_test_labelenc = init_classes_transfer.argmax(1).cpu().numpy()
form_test_labelenc = target_classes_transfer.argmax(1).cpu().numpy()
predicted_celltype_test_raw = celltype_clf_expr(test_expression_tensor)
predicted_celltype_test = celltype_clf_expr(transfer_expression_tensor)
predicted_celltype_train = None
if not annot_dataloader is None:
annot_expression = annot_dataloader.dataset.tensors[0].to(device)
predicted_celltype_train = celltype_clf_expr(annot_expression)
predicted_form_test = form_clf_expr(transfer_expression_tensor)
#predict celltypes and batch indices by model's latents
predicted_celltype_test_latents = celltype_clf_latents(test_latents)
predicted_celltype_transfer_latents = celltype_clf_latents(
transfered_latents)
predicted_form_latents = form_clf_latents(transfered_latents)
# classifier on expression
predicted_celltype_test_labelenc = predicted_celltype_test.argmax(
1).cpu().detach().numpy()
predicted_celltype_test_raw_labelenc = predicted_celltype_test_raw.argmax(
1).cpu().detach().numpy()
predicted_celltype_train_labelenc = None
if not predicted_celltype_train is None:
predicted_celltype_train_labelenc = predicted_celltype_train.argmax(
1).cpu().detach().numpy()
sameclass_mask = (init_classes_transfer.argmax(
1) == target_classes_transfer.argmax(1)).cpu().numpy().astype('bool')
celltype_train_accuracy = None
if not predicted_celltype_train_labelenc is None:
celltype_train_accuracy = (predicted_celltype_train_labelenc ==
celltype_train_labelenc).mean()
celltype_test_raw_accuracy = (predicted_celltype_test_raw_labelenc ==
np.array(celltype_test_raw_labelenc)).mean()
celltype_test_accuracy = (
predicted_celltype_test_labelenc == celltype_test_labelenc).mean()
celltype_notransfer_accuracy = (
predicted_celltype_test_labelenc == celltype_test_labelenc).mean(0)
recon_confusion = confusion_matrix(
celltype_test_labelenc[sameclass_mask],
predicted_celltype_test_labelenc[sameclass_mask])
if cfg.verbose in ('all', 'test'):
print('Confusion matrix | CELLTYPE RECONSTRUCTION')
print(recon_confusion)
from sklearn.metrics import classification_report
rep = classification_report(
celltype_test_labelenc[sameclass_mask],
predicted_celltype_test_labelenc[sameclass_mask],
output_dict=True)
df = pd.DataFrame(rep).transpose()
if cfg.verbose in ('all', 'test'):
print(df)
print('Celltype reconstruction report')
if (~sameclass_mask).any():
celltype_transfer_accuracy = (
predicted_celltype_test_labelenc[~sameclass_mask] ==
celltype_test_labelenc[~sameclass_mask]).mean(0)
transfer_confusion = confusion_matrix(
celltype_test_labelenc[~sameclass_mask],
predicted_celltype_test_labelenc[~sameclass_mask])
if cfg.verbose in ('all', 'test'):
print('Confusion matrix | CELLTYPE TRANSFER')
print(transfer_confusion)
rep = classification_report(
celltype_test_labelenc[~sameclass_mask],
predicted_celltype_test_labelenc[~sameclass_mask],
output_dict=True)
df = pd.DataFrame(rep).transpose()
if cfg.verbose in ('all', 'test'):
print('Celltype transfer report')
print(df)
pred_f_t = predicted_form_test if isinstance(
predicted_form_test, np.ndarray) else predicted_form_test.argmax(
1).cpu().detach().numpy()
rep = classification_report(form_test_labelenc[~sameclass_mask],
pred_f_t[~sameclass_mask],
output_dict=True)
df = pd.DataFrame(rep).transpose()
if cfg.verbose in ('all', 'test'):
print('Form transfer report')
print(df)
else:
celltype_transfer_accuracy = "Not enoug classes for transfering"
predicted_form_test = predicted_form_test.argmax(1).cpu().detach().numpy()
form_test_accuracy = (predicted_form_test == form_test_labelenc).mean(0)
recon_confusion = confusion_matrix(form_test_labelenc[sameclass_mask],
predicted_form_test[sameclass_mask],
normalize='all')
if cfg.verbose in ('all', 'test'):
print('Confusion matrix | FORM RECONSTRUCTED')
print(recon_confusion)
transfer_confusion = confusion_matrix(form_test_labelenc[~sameclass_mask],
pred_f_t[~sameclass_mask],
normalize='all')
if cfg.verbose in ('all', 'test'):
print('Confusion matrix | FORM TRANSFER')
print(transfer_confusion)
rep = classification_report(form_test_labelenc[~sameclass_mask],
pred_f_t[~sameclass_mask],
output_dict=True)
df = pd.DataFrame(rep).transpose()
if cfg.verbose in ('all', 'test'):
print('Form reconstructed report')
print(df)
# classifier on latents
predicted_celltype_latents_labelenc = predicted_celltype_test_latents.argmax(
1).cpu().detach().numpy()
predicted_celltype_transfer_latents_labelenc = predicted_celltype_transfer_latents.argmax(
1).cpu().detach().numpy()
#celltype_test_latents_accuracy = (predicted_celltype_latents_labelenc == celltype_test_raw_labelenc).mean()
celltype_notransfer_latents_accuracy = (
np.array(predicted_celltype_latents_labelenc) == np.array(
celltype_test_raw_labelenc)).mean(0)
if (~sameclass_mask).any():
celltype_transfer_latents_accuracy = (
predicted_celltype_transfer_latents_labelenc[~sameclass_mask] ==
celltype_test_labelenc[~sameclass_mask]).mean()
else:
celltype_transfer_latents_accuracy = "Not enoug classes for transfering"
#reconstructed means transfered to the same class
if cfg.verbose in ('all', 'test'):
print('On expression:')
print('\tCell type prediction accuracy [vae transfered data]',
celltype_test_accuracy, '< with the same class transfer')
print('\tCell type prediction accuracy [reconstructed data]',
celltype_notransfer_accuracy)
print(
'\tCell type prediction accuracy [transfered to another classes]',
celltype_transfer_accuracy)
print('\tClass prediction accuracy:', form_test_accuracy)
print('On latents:')
#print('\tCell type prediction accuracy [vae transfered data]', celltype_test_latents_accuracy,
# '< with the same class transfer')
print('\tCell type prediction accuracy [reconstructed data]',
celltype_notransfer_latents_accuracy)
print(
'\tCell type prediction accuracy [transfered to another classes]',
celltype_transfer_latents_accuracy)
#print('K neighbors purity [vae test data]:', test_purity)
#print('K neighbors purity [transfer]:', transfered_purity)
#print('Entropy batch mixing [vae test data]:', ebm_score['test'])
#print('Entropy batch mixing [transfered data]:', ebm_score['transfered'])
return {
'MSE':
float(np.mean(mse)),
'Calibration accuracy':
float((res_norm.argmin(1) == res_equal).mean()),
'Cell type prediction accuracy [train]:':
celltype_train_accuracy,
'Cell type prediction accuracy expression [test hold out data]':
celltype_test_raw_accuracy,
#'Cell type prediction accuracy latents [test hold out data]': celltype_test_latents_accuracy,
#'Cell type prediction accuracy expression [vae transfered data]': celltype_test_accuracy,
#'Cell type prediction accuracy latents [vae transfered data]': celltype_test_accuracy,
'Cell type prediction accuracy expression [reconstructed data]':
celltype_notransfer_accuracy,
'Cell type prediction accuracy latents [reconstructed data]':
celltype_notransfer_latents_accuracy,
'Cell type prediction accuracy expression [transfered to another classes]':
celltype_transfer_accuracy,
'Cell type prediction accuracy latents [transfered to another classes]':
celltype_transfer_latents_accuracy,
'Class prediction accuracy':
form_test_accuracy,
#'K neighbors purity [vae test data]': test_purity,
#'K neighbors purity [transfer]': transfered_purity,
#'Entropy batch mixing [test data]': ebm_score['test'],
#'Entropy batch mixing [transfer data]': ebm_score['transfered']
}
from torch.cuda import is_available as cuda_is_available
from pycm import ConfusionMatrix
import numpy as np
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from torch.optim.lr_scheduler import ExponentialLR
from transformers import get_linear_schedule_with_warmup
import argparse
from transformers.models.bert import BertTokenizer, BertForSequenceClassification
from nlp_utils.tokenizer import WhiteSpaceTokenizer
from nlp_utils.dataset import prepare_qqsim_train_test_org_data_loader, prepare_qqsim_predict_org_data_loader, prepare_qqsim_train_test_tensor_data_loader
from torch.nn.utils import clip_grad_norm_
from pytorch_pretrained_bert.optimization import BertAdam
from torch.nn.functional import softmax
device = device("cuda:0" if cuda_is_available() else "cpu")
#device = device("cpu")
class Config:
last_model_path = ''
pretrained_model_folder = '../user_data/pretrained/self_nezha_bert_double_sentence'
train_batch_size = 128
test_batch_size = 128
predict_batch_size = 128
epoch = 5
test_model_num = 10
log_dir = './output'
model_folder = './output'
num_labels = 2
def train(
*,
instance_file: str,
tensors_dir: str,
train_dir: str,
configs_dir: str,
model_encoder_iterations: int,
model_encoder_output_dim: int,
model_encoder_message_dim: int,
model_decoder_type: str,
model_learning_rate: float,
model_batch_size: int,
trainer_epochs: int,
trainer_eval_every: int,
trainer_limit_epochs_at: Optional[int],
trainer_train_eval_split: float,
trainer_selection_metric: str,
trainer_kept_checkpoints: int,
trainer_cuda: Optional[int],
log_level: str,
) -> None:
"""Run the training."""
Config.from_arguments(locals(),
["instance_file", "tensors_dir", "train_dir"],
"configs_dir").save(
Path(configs_dir) / "train.json")
logger = setup_logging(__name__, log_level)
tensors_dir_path = Path(tensors_dir).expanduser().resolve()
train_dir_path = Path(train_dir).expanduser().resolve()
train_dir_path.mkdir(parents=True, exist_ok=True)
with bz2_open(instance_file, "rb") as fh:
instance = pickle_load(fh)
dataset = CodRepDataset(input_dir=tensors_dir_path)
logger.info("Dataset of size %d", len(dataset))
train_length = round(0.9 * len(dataset))
eval_length = round(0.05 * len(dataset))
test_length = len(dataset) - train_length - eval_length
train_dataset, eval_dataset, test_dataset = random_split(
dataset, [train_length, eval_length, test_length])
if trainer_cuda is not None:
if not cuda_is_available():
raise RuntimeError("CUDA is not available on this system.")
device = torch_device("cuda:%d" % trainer_cuda)
else:
device = torch_device("cpu")
model = build_model(
instance=instance,
model_encoder_iterations=model_encoder_iterations,
model_encoder_output_dim=model_encoder_output_dim,
model_encoder_message_dim=model_encoder_message_dim,
model_decoder_type=model_decoder_type,
model_learning_rate=model_learning_rate,
model_batch_size=model_batch_size,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
test_dataset=test_dataset,
)
# The model needs a forward to be completely initialized.
model.training_step(instance.collate([dataset[0]]), 0)
logger.info("Configured model %s", model)
checkpoint_callback = ModelCheckpoint(
filepath=train_dir,
save_best_only=True,
verbose=True,
monitor="eval_mrr",
mode="max",
prefix="",
)
trainer = Trainer(default_save_path=train_dir,
checkpoint_callback=checkpoint_callback)
trainer.fit(model)