diff --git a/checkpoints/MNIST/bayesian/model_lenet_lrt_softplus.pt b/checkpoints/MNIST/bayesian/model_lenet_lrt_softplus.pt new file mode 100644 index 00000000..c6a639c4 Binary files /dev/null and b/checkpoints/MNIST/bayesian/model_lenet_lrt_softplus.pt differ diff --git a/get_started.py b/get_started.py new file mode 100644 index 00000000..fff13fbf --- /dev/null +++ b/get_started.py @@ -0,0 +1,14 @@ +import modal + +app = modal.App("example-get-started") + + +@app.function() +def square(x): + print("This code is running on a remote worker!") + return x**2 + + +@app.local_entrypoint() +def main(): + print("the square is", square.remote(42)) diff --git a/main_bayesian.py b/main_bayesian.py index b14c6433..20a7c79c 100755 --- a/main_bayesian.py +++ b/main_bayesian.py @@ -19,71 +19,155 @@ # CUDA settings device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") +# Reproducibility — so accuracy changes you measure come from your edits, +# not run-to-run noise +SEED = 0 +torch.manual_seed(SEED) +np.random.seed(SEED) +if torch.cuda.is_available(): + torch.cuda.manual_seed_all(SEED) + + def getModel(net_type, inputs, outputs, priors, layer_type, activation_type): - if (net_type == 'lenet'): + if net_type == 'lenet': return BBBLeNet(outputs, inputs, priors, layer_type, activation_type) - elif (net_type == 'alexnet'): + elif net_type == 'alexnet': return BBBAlexNet(outputs, inputs, priors, layer_type, activation_type) - elif (net_type == '3conv3fc'): + elif net_type == '3conv3fc': return BBB3Conv3FC(outputs, inputs, priors, layer_type, activation_type) else: - raise ValueError('Network should be either [LeNet / AlexNet / 3Conv3FC') + raise ValueError( + 'Network should be either [LeNet / AlexNet / 3Conv3FC]' + ) -def train_model(net, optimizer, criterion, trainloader, num_ens=1, beta_type=0.1, epoch=None, num_epochs=None): +def train_model( + net, + optimizer, + criterion, + trainloader, + num_ens=1, + beta_type=0.1, + epoch=None, + num_epochs=None, + grad_clip_norm=5.0, +): net.train() + training_loss = 0.0 accs = [] kl_list = [] + for i, (inputs, labels) in enumerate(trainloader, 1): optimizer.zero_grad() inputs, labels = inputs.to(device), labels.to(device) - outputs = torch.zeros(inputs.shape[0], net.num_classes, num_ens).to(device) + + outputs = torch.zeros( + inputs.shape[0], + net.num_classes, + num_ens + ).to(device) kl = 0.0 + for j in range(num_ens): net_out, _kl = net(inputs) kl += _kl outputs[:, :, j] = F.log_softmax(net_out, dim=1) - + kl = kl / num_ens kl_list.append(kl.item()) + log_outputs = utils.logmeanexp(outputs, dim=2) - beta = metrics.get_beta(i-1, len(trainloader), beta_type, epoch, num_epochs) + beta = metrics.get_beta( + i - 1, + len(trainloader), + beta_type, + epoch, + num_epochs + ) + loss = criterion(log_outputs, labels, kl, beta) + loss.backward() + + # Gradient clipping — BBB layers sample noise every forward pass, + # which makes gradients noisier than a plain CNN; clipping + # stabilizes training and helps final accuracy. + torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=grad_clip_norm) + optimizer.step() accs.append(metrics.acc(log_outputs.data, labels)) - training_loss += loss.cpu().data.numpy() - return training_loss/len(trainloader), np.mean(accs), np.mean(kl_list) + training_loss += loss.item() + + return ( + training_loss / len(trainloader), + np.mean(accs), + np.mean(kl_list) + ) -def validate_model(net, criterion, validloader, num_ens=1, beta_type=0.1, epoch=None, num_epochs=None): +def validate_model( + net, + criterion, + validloader, + num_ens=1, + beta_type=0.1, + epoch=None, + num_epochs=None +): """Calculate ensemble accuracy and NLL Loss""" - net.train() + + net.eval() + valid_loss = 0.0 accs = [] - for i, (inputs, labels) in enumerate(validloader): - inputs, labels = inputs.to(device), labels.to(device) - outputs = torch.zeros(inputs.shape[0], net.num_classes, num_ens).to(device) - kl = 0.0 - for j in range(num_ens): - net_out, _kl = net(inputs) - kl += _kl - outputs[:, :, j] = F.log_softmax(net_out, dim=1).data + with torch.no_grad(): - log_outputs = utils.logmeanexp(outputs, dim=2) + for i, (inputs, labels) in enumerate(validloader): + + inputs, labels = inputs.to(device), labels.to(device) + + outputs = torch.zeros( + inputs.shape[0], + net.num_classes, + num_ens + ).to(device) + + kl = 0.0 - beta = metrics.get_beta(i-1, len(validloader), beta_type, epoch, num_epochs) - valid_loss += criterion(log_outputs, labels, kl, beta).item() - accs.append(metrics.acc(log_outputs, labels)) + for j in range(num_ens): + net_out, _kl = net(inputs) + kl += _kl + outputs[:, :, j] = F.log_softmax(net_out, dim=1) - return valid_loss/len(validloader), np.mean(accs) + kl = kl / num_ens + + log_outputs = utils.logmeanexp(outputs, dim=2) + + beta = metrics.get_beta( + i - 1, + len(validloader), + beta_type, + epoch, + num_epochs + ) + + valid_loss += criterion( + log_outputs, + labels, + kl, + beta + ).item() + + accs.append(metrics.acc(log_outputs, labels)) + + return valid_loss / len(validloader), np.mean(accs) def run(dataset, net_type): @@ -103,40 +187,158 @@ def run(dataset, net_type): beta_type = cfg.beta_type trainset, testset, inputs, outputs = data.getDataset(dataset) + train_loader, valid_loader, test_loader = data.getDataloader( - trainset, testset, valid_size, batch_size, num_workers) - net = getModel(net_type, inputs, outputs, priors, layer_type, activation_type).to(device) + trainset, + testset, + valid_size, + batch_size, + num_workers + ) + + net = getModel( + net_type, + inputs, + outputs, + priors, + layer_type, + activation_type + ).to(device) ckpt_dir = f'checkpoints/{dataset}/bayesian' - ckpt_name = f'checkpoints/{dataset}/bayesian/model_{net_type}_{layer_type}_{activation_type}.pt' + ckpt_name = ( + f'checkpoints/{dataset}/bayesian/' + f'model_{net_type}_{layer_type}_{activation_type}.pt' + ) - if not os.path.exists(ckpt_dir): - os.makedirs(ckpt_dir, exist_ok=True) + os.makedirs(ckpt_dir, exist_ok=True) criterion = metrics.ELBO(len(trainset)).to(device) - optimizer = Adam(net.parameters(), lr=lr_start) - lr_sched = lr_scheduler.ReduceLROnPlateau(optimizer, patience=6, verbose=True) - valid_loss_max = np.Inf - for epoch in range(n_epochs): # loop over the dataset multiple times - train_loss, train_acc, train_kl = train_model(net, optimizer, criterion, train_loader, num_ens=train_ens, beta_type=beta_type, epoch=epoch, num_epochs=n_epochs) - valid_loss, valid_acc = validate_model(net, criterion, valid_loader, num_ens=valid_ens, beta_type=beta_type, epoch=epoch, num_epochs=n_epochs) + # weight_decay added — the KL term already regularizes toward the + # prior, so keep this small relative to a plain CNN. + optimizer = Adam(net.parameters(), lr=lr_start, weight_decay=1e-5) + + lr_sched = lr_scheduler.ReduceLROnPlateau( + optimizer, + patience=6 + ) + + # Track best VALIDATION ACCURACY, not lowest validation loss. + # valid_loss = NLL + beta * KL, and beta changes across epochs + # (Blundell schedule), so it isn't a stable, comparable quantity + # epoch-to-epoch. Accuracy is what we actually care about. + best_valid_acc = 0.0 + + for epoch in range(n_epochs): + + train_loss, train_acc, train_kl = train_model( + net, + optimizer, + criterion, + train_loader, + num_ens=train_ens, + beta_type=beta_type, + epoch=epoch, + num_epochs=n_epochs + ) + + # More MC samples at validation time = lower-variance, generally + # more accurate estimate of the model's true predictive + # performance, at the cost of extra forward passes (no retraining + # needed). Bump this if you can afford the compute. + eval_ens = max(valid_ens, 10) + + valid_loss, valid_acc = validate_model( + net, + criterion, + valid_loader, + num_ens=eval_ens, + beta_type=beta_type, + epoch=epoch, + num_epochs=n_epochs + ) + + old_lr = optimizer.param_groups[0]['lr'] + lr_sched.step(valid_loss) - print('Epoch: {} \tTraining Loss: {:.4f} \tTraining Accuracy: {:.4f} \tValidation Loss: {:.4f} \tValidation Accuracy: {:.4f} \ttrain_kl_div: {:.4f}'.format( - epoch, train_loss, train_acc, valid_loss, valid_acc, train_kl)) + new_lr = optimizer.param_groups[0]['lr'] + + if old_lr != new_lr: + print( + f"Learning rate reduced from " + f"{old_lr:.6e} to {new_lr:.6e}" + ) + + print( + 'Epoch: {} \tTraining Loss: {:.4f} ' + '\tTraining Accuracy: {:.4f} ' + '\tValidation Loss: {:.4f} ' + '\tValidation Accuracy: {:.4f} ' + '\ttrain_kl_div: {:.4f}'.format( + epoch, + train_loss, + train_acc, + valid_loss, + valid_acc, + train_kl + ) + ) + + if valid_acc >= best_valid_acc: + print( + 'Validation accuracy increased ' + '({:.4f} --> {:.4f}). Saving model ...'.format( + best_valid_acc, + valid_acc + ) + ) - # save model if validation accuracy has increased - if valid_loss <= valid_loss_max: - print('Validation loss decreased ({:.6f} --> {:.6f}). Saving model ...'.format( - valid_loss_max, valid_loss)) torch.save(net.state_dict(), ckpt_name) - valid_loss_max = valid_loss + best_valid_acc = valid_acc + + # Final honest number: reload the best checkpoint and evaluate on the + # held-out test set (previously the script never touched test_loader, + # so you had no unbiased accuracy figure at all). + net.load_state_dict(torch.load(ckpt_name, map_location=device)) + test_loss, test_acc = validate_model( + net, + criterion, + test_loader, + num_ens=max(valid_ens, 20), + beta_type=beta_type, + epoch=n_epochs, + num_epochs=n_epochs + ) + print( + 'Best Validation Accuracy: {:.4f}' + '\tTest Loss: {:.4f} \tTest Accuracy: {:.4f}'.format( + best_valid_acc, test_loss, test_acc + ) + ) + if __name__ == '__main__': - parser = argparse.ArgumentParser(description = "PyTorch Bayesian Model Training") - parser.add_argument('--net_type', default='lenet', type=str, help='model') - parser.add_argument('--dataset', default='MNIST', type=str, help='dataset = [MNIST/CIFAR10/CIFAR100]') + + parser = argparse.ArgumentParser( + description="PyTorch Bayesian Model Training" + ) + + parser.add_argument( + '--net_type', + default='lenet', + type=str, + help='model' + ) + + parser.add_argument( + '--dataset', + default='MNIST', + type=str, + help='dataset = [MNIST/CIFAR10/CIFAR100]' + ) + args = parser.parse_args() - run(args.dataset, args.net_type) + run(args.dataset, args.net_type) \ No newline at end of file diff --git a/modal_train.py b/modal_train.py new file mode 100644 index 00000000..76fd10f9 --- /dev/null +++ b/modal_train.py @@ -0,0 +1,95 @@ +import subprocess +import modal + +# Create the Modal application +app = modal.App("bayesian-cnn") + +# Persistent storage for trained checkpoints. Without this, anything written +# inside the container's filesystem (including checkpoints/) is destroyed the +# moment the function returns — that's why your last run "vanished". +volume = modal.Volume.from_name("bayesiancnn-checkpoints", create_if_missing=True) + +CHECKPOINT_DIR = "/root/project/checkpoints" + +# Build the execution environment +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch", + "torchvision", + "numpy", + "scipy", + "matplotlib", + "tqdm", + "tensorboard", + "Pillow", + ) + .add_local_dir( + ".", + remote_path="/root/project", + ignore=[ + ".git", + "venv", + "__pycache__", + ".pytest_cache", + ".idea", + ".vscode", + ".github", + "checkpoints", + "*.pyc", + "*.pyo", + "*.pyd", + ".DS_Store", + ], + ) +) + + +@app.function( + image=image, + gpu="A10G", + cpu=4, + memory=16384, + timeout=60 * 60 * 12, # 12 hours + volumes={CHECKPOINT_DIR: volume}, +) +def train(dataset: str = "CIFAR10", net_type: str = "3conv3fc"): + + # Confirms the container actually has a GPU attached before you wait + # hours for a run — if this fails/prints nothing useful, the job is + # running on CPU regardless of gpu="A10G" being set. + subprocess.run(["nvidia-smi"], cwd="/root/project") + + print(f"Starting Bayesian CNN training: dataset={dataset}, net_type={net_type}") + + # Popen + streaming instead of subprocess.run(capture_output=True): + # capture_output buffers everything and only prints it after the + # process exits, so a multi-hour job would show nothing in the Modal + # dashboard until it was already finished (or silently hung). + process = subprocess.Popen( + ["python", "main_bayesian.py", "--dataset", dataset, "--net_type", net_type], + cwd="/root/project", + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + bufsize=1, + ) + for line in process.stdout: + print(line, end="") + + process.wait() + + if process.returncode != 0: + raise RuntimeError( + f"Training failed with exit code {process.returncode}" + ) + + # Persist everything written to CHECKPOINT_DIR during this run back to + # the Modal Volume so it survives after the container is torn down. + volume.commit() + print(f"Checkpoints committed to volume 'bayesiancnn-checkpoints'.") + + +@app.local_entrypoint() +def main(dataset: str = "CIFAR10", net_type: str = "3conv3fc"): + train.remote(dataset=dataset, net_type=net_type) \ No newline at end of file diff --git a/mycontributions/Read-notes.md b/mycontributions/Read-notes.md new file mode 100644 index 00000000..824921f9 --- /dev/null +++ b/mycontributions/Read-notes.md @@ -0,0 +1 @@ +I improved how the main_bayesian.py handles data on lower pc diff --git a/mycontributions/get_started.py b/mycontributions/get_started.py new file mode 100644 index 00000000..fff13fbf --- /dev/null +++ b/mycontributions/get_started.py @@ -0,0 +1,14 @@ +import modal + +app = modal.App("example-get-started") + + +@app.function() +def square(x): + print("This code is running on a remote worker!") + return x**2 + + +@app.local_entrypoint() +def main(): + print("the square is", square.remote(42))