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):