# Plot the training images.
sample_batch = next(iter(dataloader))
plt.figure(figsize=(10, 10))
plt.axis("off")
plt.imshow(
np.transpose(
vutils.make_grid(sample_batch[0].to(device)[:100],
nrow=10,
padding=2,
normalize=True).cpu(), (1, 2, 0)))
plt.savefig('Training Images {}'.format(params['dataset']))
plt.close('all')
# Initialise the network.
netG = Generator().to(device)
netG.apply(weights_init)
print(netG)
discriminator = Discriminator().to(device)
discriminator.apply(weights_init)
print(discriminator)
netD = DHead().to(device)
netD.apply(weights_init)
print(netD)
netQ = QHead().to(device)
netQ.apply(weights_init)
print(netQ)
import matplotlib.pyplot as plt
from config import params as params_config
from models.mnist_model import Generator
device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")
state_dict = torch.load(params_config['load_path'])
params = state_dict['params']
anomaly_label = params_config['anomaly_label']
filename = params_config['filename']
dis_c_dim = params['dis_c_dim']
# Create the generator network.
num_z_c = params['num_z'] + params['num_dis_c'] * params['dis_c_dim'] + params[
'num_con_c']
netG = Generator(num_z_c).to(device)
# Load the trained generator weights.
netG.load_state_dict(state_dict['netG'])
# print(netG)
c = np.linspace(-2, 2, params['dis_c_dim']).reshape(1, -1)
c = np.repeat(c, dis_c_dim, 0).reshape(-1, 1)
c = torch.from_numpy(c).float().to(device)
c = c.view(-1, 1, 1, 1)
dis_c_dim_squared = dis_c_dim * dis_c_dim
zeros = torch.zeros(dis_c_dim_squared, 1, 1, 1, device=device)
idx = np.arange(params['dis_c_dim']).repeat(dis_c_dim)
dis_c = torch.zeros(dis_c_dim_squared,
params['dis_c_dim'],