def __init__(self, context: PyTorchTrialContext) -> None:
# Read configuration
self.context = context
self.data_config = self.context.get_data_config()
# Create Tensorboard logger
self.logger = TorchWriter()
# Create tokinizer based on the predefient vocabulary
self.tokenizer = BertTokenizer(self.data_config["voc_path"], do_lower_case=False)
# Label Encoder
if self.context.get_hparam("reduce_to_binary_problem"):
class_num = 2
else:
class_num = 6
# Initialize model and wrap it in the determined api
model = ProtTransClassification(self.data_config["pretrained_path"],
class_num=class_num,
classification_feature=self.context.get_hparam("classification_feature"),
dropout=self.context.get_hparam("classification_dropout"),
freeze_bert=self.context.get_hparam("bert_freeze"))
optimizer = Lamb([{"params": model.wordencoding.parameters(), "lr": self.context.get_hparam("bert_lr")},
{"params": model.classification.parameters()}], lr=self.context.get_hparam("classification_lr"))
self.model = self.context.wrap_model(model)
self.optimizer = self.context.wrap_optimizer(optimizer)
def __init__(self, context: PyTorchTrialContext) -> None:
self.context = context
self.logger = TorchWriter()
# Create a unique download directory for each rank so they don't overwrite each other.
self.download_directory = f"/tmp/data-rank{self.context.distributed.get_rank()}"
self.data_downloaded = False
# Initialize the models.
mnist_shape = (1, 28, 28)
self.generator = self.context.wrap_model(Generator(latent_dim=self.context.get_hparam("latent_dim"),
img_shape=mnist_shape))
self.discriminator = self.context.wrap_model(Discriminator(img_shape=mnist_shape))
# Initialize the optimizers and learning rate scheduler.
lr = self.context.get_hparam("lr")
b1 = self.context.get_hparam("b1")
b2 = self.context.get_hparam("b2")
self.opt_g = self.context.wrap_optimizer(torch.optim.Adam(self.generator.parameters(),
lr=lr, betas=(b1, b2)))
self.opt_d = self.context.wrap_optimizer(torch.optim.Adam(self.discriminator.parameters(),
lr=lr, betas=(b1, b2)))
self.lr_g = self.context.wrap_lr_scheduler(
lr_scheduler=LambdaLR(self.opt_g, lr_lambda=lambda epoch: 0.95 ** epoch),
step_mode=LRScheduler.StepMode.STEP_EVERY_EPOCH,
)
def create_metric_writer(
cls: Type["PyTorchTrialController"],
) -> tensorboard.BatchMetricWriter:
from determined.tensorboard.metric_writers.pytorch import TorchWriter
writer = TorchWriter()
return tensorboard.BatchMetricWriter(writer)
def __init__(self, context: DeepSpeedTrialContext) -> None:
self.context = context
self.exp_config = self.context.get_experiment_config()
self.args = AttrMap(self.context.get_hparams())
# Initalize and get arguments, timers, and Tensorboard writer.
try:
self.neox_args = get_neox_args(self.context)
except:
traceback.print_exc()
raise InvalidHP("Could not parse neox_args.")
self.wrapped_writer = TorchWriter()
self.neox_args.tensorboard_writer = self.wrapped_writer.writer
self.neox_args.configure_distributed_args()
# The tokenizer needs to be built before model initialization in order to set the
# required padded_vocab_size argument.
self.neox_args.build_tokenizer()
megatron_train.initialize_megatron(neox_args=self.neox_args)
self.timers = megatron_utils.Timers(
use_wandb=False,
tensorboard_writer=self.neox_args.tensorboard_writer)
# Model, optimizer, and learning rate.
self.timers("model and optimizer").start()
(
model,
self.optimizer,
self.lr_scheduler,
) = megatron_train.setup_model_and_optimizer(neox_args=self.neox_args)
self.model = self.context.wrap_model_engine(model)
self.timers("model and optimizer").stop()
# Print setup timing.
megatron_utils.print_rank_0("done with setups ...")
self.timers.log(["model and optimizer"])
megatron_utils.print_rank_0("training ...")
# For tracking.
if not self.args.search_world_size:
self.reducer = self.context.wrap_reducer(LMReducers(
self.neox_args),
for_training=False,
for_validation=True)
self.report_memory_flag = True
self.total_train_loss_dict = {}
self.total_val_loss_dict = {}
self.tflops = 0
self.reported_flops = False
self.overflow_monitor = megatron_utils.OverflowMonitor(self.optimizer)
self.noise_scale_logger = megatron_utils.get_noise_scale_logger(
self.neox_args)
self.timers("interval time").start()
def __init__(self, context: DeepSpeedTrialContext) -> None:
self.context = context
self.hparams = AttrDict(self.context.get_hparams())
self.data_config = AttrDict(self.context.get_data_config())
self.logger = TorchWriter()
num_channels = data.CHANNELS_BY_DATASET[self.data_config.dataset]
gen_net = Generator(
self.hparams.generator_width_base, num_channels, self.hparams.noise_length
)
gen_net.apply(weights_init)
disc_net = Discriminator(self.hparams.discriminator_width_base, num_channels)
disc_net.apply(weights_init)
gen_parameters = filter(lambda p: p.requires_grad, gen_net.parameters())
disc_parameters = filter(lambda p: p.requires_grad, disc_net.parameters())
ds_config = overwrite_deepspeed_config(
self.hparams.deepspeed_config, self.hparams.get("overwrite_deepspeed_args", {})
)
generator, _, _, _ = deepspeed.initialize(
model=gen_net, model_parameters=gen_parameters, config=ds_config
)
discriminator, _, _, _ = deepspeed.initialize(
model=disc_net, model_parameters=disc_parameters, config=ds_config
)
self.generator = self.context.wrap_model_engine(generator)
self.discriminator = self.context.wrap_model_engine(discriminator)
self.fixed_noise = self.context.to_device(
torch.randn(
self.context.train_micro_batch_size_per_gpu, self.hparams.noise_length, 1, 1
)
)
self.criterion = nn.BCELoss()
# TODO: Test fp16
self.fp16 = generator.fp16_enabled()
self.gradient_accumulation_steps = generator.gradient_accumulation_steps()
# Manually perform gradient accumulation.
if self.gradient_accumulation_steps > 1:
logging.info("Disabling automatic gradient accumulation.")
self.context.disable_auto_grad_accumulation()
def create_metric_writer(
cls: Type["DeepSpeedTrialController"],
) -> tensorboard.BatchMetricWriter:
writer = TorchWriter()
return tensorboard.BatchMetricWriter(writer)