def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(device) seed = 42 out_dir = './logs' if not os.path.exists(out_dir): os.mkdir(out_dir) checkpoints_dir = "./checkpoints" if not os.path.exists(checkpoints_dir): os.mkdir(out_dir) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) model = VAE(z_dim=512) model.load_state_dict(torch.load("./checkpoints/500.pth")) model = model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=5e-4) test_loader = return_MVTecAD_loader( image_dir="./mvtec_anomaly_detection/grid/test/metal_contamination/", batch_size=10, train=False) #eval(model=model,test_loader=test_loader,device=device) EBM(model, test_loader, device)
print("Random Seed: ", opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) # specify the gpu id if using only 1 gpu if opt.ngpu == 1: os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.gpu_id) if opt.cuda: torch.cuda.manual_seed_all(opt.manualSeed) cudnn.benchmark = True # model netVAE = VAE('test') if opt.netVAE != '': netVAE.load_state_dict(torch.load(opt.netVAE)['state_dict']) print('loading model ....') netVAE.eval() print(netVAE) if opt.cuda: netVAE.cuda() print("Saving random generation......") decoder = netVAE.decoder # method 1 noise_1 = torch.randn(1024, 5, 5) noise_2 = torch.randn(1024, 5, 5) noise = [] for i in range(32):