def test():
# set initial config
set_config()
# load config from wandb
logger.load_last_wandb_config()
# load last state dict from wandb
model = setup_model()
ckpt = logger.load_last_wandb_checkpoint()
# load state dict
ms = model.module if cfg.GRADCAM.NUM_GPUS > 1 else model
ms.load_state_dict(ckpt['model_state'])
# get test data
gradcam_loader = loader.construct_gradcam_loader()
# seletect module where you want gradcam
feature_module = model.module.stem if cfg.GRADCAM.NUM_GPUS > 1 else model.stem
print(feature_module)
gradcam = GradCAM(model, feature_module, ["conv"])
for cur_iter, (inputs, labels) in enumerate(gradcam_loader):
# Transfer the data to the current GPU device
inputs, labels = inputs.cuda(), labels.cuda(non_blocking=True)
mask = gradcam(inputs)
img = get_image_from_tensor(inputs)
show_cam_on_image(img, mask)
def test_model():
# set initial config
set_config()
# load config from wandb
logger.load_last_wandb_config()
# load last model from wandb
model = setup_model()
logger.load_last_wandb_checkpoint(model)
# get test data
test_loader = loader.construct_test_loader()
def setup_env():
"""Sets up environment for training or testing."""
# load environment variable rom file
load_dotenv()
# set config variable
set_config()
# init wandb logger
logger.init_logger()
# save cfg to wandb
logger.save_config()
torch.cuda.empty_cache()
np.random.seed(cfg.RANDOM_SEED)
torch.manual_seed(cfg.RANDOM_SEED)
def test_gradcam():
# set initial config
set_config()
# load config from wandb
logger.load_last_wandb_config()
# load last model from wandb
model = setup_model()
logger.load_last_wandb_checkpoint(model)
# get test data
test_loader = loader.construct_test_loader()
gradcam = GradCAM(model, model.head, ["fc"])
for inputs, labels in test_loader:
inputs, labels = inputs.cuda(), labels.cuda(non_blocking=True)
# Compute the predictions
grad, preds = gradcam(inputs)
print(grad[0].shape)
def save_checkpoint(model, optim, epoch, wandb_path = True):
sd = model.module.state_dict() if cfg.NUM_GPUS > 1 else model.state_dict()
ckpt = {
"epoch": epoch,
"model_state": sd,
"optim_state": optim.state_dict()
}
if wandb_path:
path = get_wandb_checkpoint()
else:
path = get_checkpoint()
# save to local path
torch.save(ckpt, path)
# save to wandb
# wandb.save(path)
print(f"Save to {path}")
def load_wandb_last_checkpoint():
run = get_wandb_last_run()
model = get_model()
if __name__ =='__main__':
set_config()
load_last_wandb_checkpoint(5)