Loading a Model

[jprezioso]

Thank you for sharing this code base! I have a trained BNN and am having some trouble defining my trained model after loading up the state dictionary for the trained model. Specifically, I am getting an error when calling the getModel() function. There is some code below which reproduces the error, which seems to be an issue with creating the path to save the means and variances. If I turn off the saving feature, I get another TypeError related to the model and layer classes. Can you provide an example of loading up a trained model or some additional insight? Thanks again!

--

from future import print_function

import os import argparse

import torch import numpy as np from torch.optim import Adam from torch.nn import functional as F

import data import utils import metrics import config_bayesian as cfg from models.BayesianModels.Bayesian3Conv3FC import BBB3Conv3FC from models.BayesianModels.BayesianAlexNet import BBBAlexNet from models.BayesianModels.BayesianLeNet import BBBLeNet from models.BayesianModels.BayesianFeedForward import BBBFeedForward from main_bayesian import getModel

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

net_type = 'lenet' trainset, testset, inputs, outputs = data.getDataset('MNIST') net = getModel(net_type, inputs, outputs).to(device)

[Piyush-555]

Hi @jprezioso! Looks like you're using the previous version of the code. The repository is going through several changes lately.

Here's a code for loading an already trained model. Try this with the updated code.

import torch

import data
import metrics
import config_bayesian as cfg
from main_bayesian import getModel, validate_model

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

dataset = 'MNIST'
net_type = 'lenet'

weight_path = f'checkpoints/MNIST/bayesian/model_lenet_lrt_softplus.pt'  # Your weights

trainset, testset, inputs, outputs = data.getDataset(dataset)

train_loader, valid_loader, test_loader = data.getDataloader(
        trainset, testset, cfg.valid_size, cfg.batch_size, cfg.num_workers)

net = getModel(net_type, inputs, outputs, None, cfg.layer_type, cfg.activation_type)  # priors could be anything, it's going to be replaced anyway
net.load_state_dict(torch.load(weight_path))
net = net.to(device)

# Here, you got your pre-trained model. Now, just validating..
criterion = metrics.ELBO(len(trainset)).to(device)
valid_loss, valid_acc = validate_model(net, criterion, valid_loader, num_ens=cfg.valid_ens, beta_type=cfg.beta_type)

print('Validation Loss: {:.4f} \tValidation Accuracy: {:.4f}'.format(valid_loss, valid_acc))

Hope this helps!

[jprezioso]

Yes, thank you!