def __init__(
self,
config: Config,
serialization_dir: str,
gpu_ids: List[int] = [0],
cpu_workers: int = 0,
):
self._C = config
if self._C.PHASE != "question_coding":
raise ValueError(
f"Trying to initialize a QuestionCodingTrainer, expected config PHASE to be "
f"question_coding, found {self._C.PHASE}")
# Initialize dataloader and model.
dataset = QuestionCodingDataset(
self._C.DATA.TRAIN_TOKENS,
num_supervision=self._C.SUPERVISION,
supervision_question_max_length=self._C.
SUPERVISION_QUESTION_MAX_LENGTH,
)
sampler = SupervisionWeightedRandomSampler(dataset)
dataloader = DataLoader(dataset,
batch_size=self._C.OPTIM.BATCH_SIZE,
sampler=sampler,
num_workers=cpu_workers)
program_generator = ProgramGenerator.from_config(self._C)
question_reconstructor = QuestionReconstructor.from_config(self._C)
super().__init__(
config=config,
dataloader=dataloader,
models={
"program_generator": program_generator,
"question_reconstructor": question_reconstructor,
},
serialization_dir=serialization_dir,
gpu_ids=gpu_ids,
)
# These will be a part of `self._models`, keep these handles for convenience.
self._program_generator = self._models["program_generator"]
self._question_reconstructor = self._models["question_reconstructor"]
# Load program prior from checkpoint, this will be frozen during question coding.
self._program_prior = ProgramPrior.from_config(self._C).to(
self._device)
self._program_prior.load_state_dict(
torch.load(self._C.CHECKPOINTS.PROGRAM_PRIOR)["program_prior"])
self._program_prior.eval()
# Instantiate an elbo module to compute evidence lower bound during `_do_iteration`.
self._elbo = QuestionCodingElbo(
program_generator=self._program_generator,
question_reconstructor=self._question_reconstructor,
program_prior=self._program_prior,
beta=self._C.BETA,
baseline_decay=self._C.DELTA,
)
def __init__(
self,
config: Config,
models: Dict[str, Type[nn.Module]],
gpu_ids: List[int] = [0],
cpu_workers: int = 0,
):
self._C = config
if self._C.PHASE != "module_training":
raise ValueError(
f"Trying to initialize a ModuleTrainingEvaluator, expected config PHASE to be "
f"module_training, found {self._C.PHASE}"
)
# Initialize dataloader and model.
dataset = ModuleTrainingDataset(
self._C.DATA.VAL_TOKENS, self._C.DATA.VAL_FEATURES, in_memory=False
)
dataloader = DataLoader(
dataset, batch_size=self._C.OPTIM.BATCH_SIZE, num_workers=cpu_workers
)
super().__init__(config=config, dataloader=dataloader, models=models, gpu_ids=gpu_ids)
# This will be a part of `self._models`, keep this handle for convenience.
self._nmn = self._models["nmn"]
# Load program generator from checkpoint, this will be frozen during module training.
self._program_generator = ProgramGenerator.from_config(self._C).to(self._device)
CheckpointManager(program_generator=self._program_generator).load(
self._C.CHECKPOINTS.QUESTION_CODING
)
self._program_generator.eval()
def __init__(
self,
config: Config,
serialization_dir: str,
gpu_ids: List[int] = [0],
cpu_workers: int = 0,
):
self._C = config
if self._C.PHASE != "module_training":
raise ValueError(
f"Trying to initialize a ModuleTrainingTrainer, expected config PHASE to be "
f"module_training, found {self._C.PHASE}"
)
# Initialize dataloader and model.
dataset = ModuleTrainingDataset(
self._C.DATA.TRAIN_TOKENS, self._C.DATA.TRAIN_FEATURES, in_memory=False
)
dataloader = DataLoader(
dataset, batch_size=self._C.OPTIM.BATCH_SIZE, num_workers=cpu_workers
)
nmn = NeuralModuleNetwork.from_config(self._C)
super().__init__(
config=config,
dataloader=dataloader,
models={"nmn": nmn},
serialization_dir=serialization_dir,
gpu_ids=gpu_ids,
)
# This will be a part of `self._models`, keep this handle for convenience.
self._nmn = self._models["nmn"]
# Load program generator from checkpoint, this will be frozen during module training.
self._program_generator = ProgramGenerator.from_config(self._C).to(self._device)
self._program_generator.load_state_dict(
torch.load(self._C.CHECKPOINTS.QUESTION_CODING)["program_generator"]
)
self._program_generator.eval()
# For reproducibility - refer https://pytorch.org/docs/stable/notes/randomness.html
# These five lines control all the major sources of randomness.
np.random.seed(_C.RANDOM_SEED)
torch.manual_seed(_C.RANDOM_SEED)
torch.cuda.manual_seed_all(_C.RANDOM_SEED)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
# --------------------------------------------------------------------------------------------
# INSTANTIATE DATALOADER AND MODELS
# --------------------------------------------------------------------------------------------
dataset = JointTrainingDataset(_C.DATA.TEST_TOKENS, _C.DATA.TEST_FEATURES)
dataloader = DataLoader(dataset, batch_size=_C.OPTIM.BATCH_SIZE, num_workers=_A.cpu_workers)
program_generator = ProgramGenerator.from_config(_C).to(device)
nmn = NeuralModuleNetwork.from_config(_C).to(device)
CheckpointManager(program_generator=program_generator, nmn=nmn).load(_A.checkpoint_path)
program_generator.eval()
nmn.eval()
# To convert answer tokens to answer strings.
vocabulary = Vocabulary.from_files(_C.DATA.VOCABULARY)
predictions: List[Dict[str, Union[int, str]]] = []
for batch in tqdm(dataloader):
for key in batch:
batch[key] = batch[key].to(device)
def __init__(
self,
config: Config,
serialization_dir: str,
gpu_ids: List[int] = [0],
cpu_workers: int = 0,
):
self._C = config
if self._C.PHASE != "joint_training":
raise ValueError(
f"Trying to initialize a JointTrainingTrainer, expected config PHASE to be "
f"joint_training, found {self._C.PHASE}"
)
# Initialize dataloader and model.
dataset = JointTrainingDataset(
self._C.DATA.TRAIN_TOKENS,
self._C.DATA.TRAIN_FEATURES,
num_supervision=self._C.SUPERVISION,
supervision_question_max_length=self._C.SUPERVISION_QUESTION_MAX_LENGTH,
)
sampler = SupervisionWeightedRandomSampler(dataset)
dataloader = DataLoader(
dataset, batch_size=self._C.OPTIM.BATCH_SIZE, sampler=sampler, num_workers=cpu_workers
)
program_generator = ProgramGenerator.from_config(self._C)
question_reconstructor = QuestionReconstructor.from_config(self._C)
nmn = NeuralModuleNetwork.from_config(self._C)
# Load checkpoints from question_coding and module_training phases.
CheckpointManager(
program_generator=program_generator, question_reconstructor=question_reconstructor
).load(self._C.CHECKPOINTS.QUESTION_CODING)
CheckpointManager(nmn=nmn).load(self._C.CHECKPOINTS.MODULE_TRAINING)
super().__init__(
config=config,
dataloader=dataloader,
models={
"program_generator": program_generator,
"question_reconstructor": question_reconstructor,
"nmn": nmn,
},
serialization_dir=serialization_dir,
gpu_ids=gpu_ids,
)
# These will be a part of `self._models`, keep these handles for convenience.
self._program_generator = self._models["program_generator"]
self._question_reconstructor = self._models["question_reconstructor"]
self._nmn = self._models["nmn"]
# Load program prior from checkpoint, this will be frozen during question coding.
self._program_prior = ProgramPrior.from_config(self._C).to(self._device)
CheckpointManager(program_prior=self._program_prior).load(
self._C.CHECKPOINTS.PROGRAM_PRIOR
)
self._program_prior.eval()
# Instantiate an elbo module to compute evidence lower bound during `_do_iteration`.
self._elbo = JointTrainingElbo(
program_generator=self._program_generator,
question_reconstructor=self._question_reconstructor,
nmn=self._nmn,
program_prior=self._program_prior,
beta=self._C.BETA,
gamma=self._C.GAMMA,
baseline_decay=self._C.DELTA,
objective=self._C.OBJECTIVE,
)