def __init__(self, z_dim=50, hidden_dim=400, use_cuda=False):
super(VAE, self).__init__()
# create the encoder and decoder networks
self.encoder = Encoder(z_dim, hidden_dim)
self.decoder = Decoder(z_dim, hidden_dim)
if use_cuda:
# calling cuda() here will put all the parameters of
# the encoder and decoder networks into gpu memory
self.cuda()
self.use_cuda = use_cuda
self.z_dim = z_dim
class VAE(nn.Module):
# by default our latent space is 50-dimensional
# and we use 400 hidden units
def __init__(self, z_dim=50, hidden_dim=400, use_cuda=False):
super(VAE, self).__init__()
# create the encoder and decoder networks
self.encoder = Encoder(z_dim, hidden_dim)
self.decoder = Decoder(z_dim, hidden_dim)
if use_cuda:
# calling cuda() here will put all the parameters of
# the encoder and decoder networks into gpu memory
self.cuda()
self.use_cuda = use_cuda
self.z_dim = z_dim
# define the model p(x|z)p(z)
def model(self, x):
# register PyTorch module `decoder` with Pyro
pyro.module("decoder", self.decoder)
with pyro.plate("data", x.shape[0]):
# setup hyperparameters for prior p(z)
z_loc = x.new_zeros(torch.Size((x.shape[0], self.z_dim)))
z_scale = x.new_ones(torch.Size((x.shape[0], self.z_dim)))
# sample from prior (value will be sampled by guide when computing the ELBO)
z = pyro.sample("latent", dist.Normal(z_loc, z_scale).to_event(1))
# decode the latent code z
loc, k = self.decoder.forward(z)
# score against actual images
pyro.sample("obs", dist.VonMises(loc, k).to_event(1), obs=x)
# define the guide (i.e. variational distribution) q(z|x)
def guide(self, x):
# register PyTorch module `encoder` with Pyro
pyro.module("encoder", self.encoder)
set_trace()
with pyro.plate("data", x.shape[0]):
# use the encoder to get the parameters used to define q(z|x)
z_loc, z_scale = self.encoder.forward(x)
# sample the latent code z
pyro.sample("latent", dist.Normal(z_loc, z_scale).to_event(1))
# define a helper function for reconstructing images
def reconstruct_img(self, x):
# encode image x
z_loc, z_scale = self.encoder(x)
# sample in latent space
z = dist.Normal(z_loc, z_scale).sample()
# decode the image (note we don't sample in image space)
loc = self.decoder(z)
return loc
netG.apply(weights_init)
pass
else:
try:
pretrained_model = torch.load('./Generators/' + args.model_id + '.pt')
try:
netG.load_state_dict(pretrained_model.state_dict())
except:
netG.load_state_dict(pretrained_model)
except:
print('G weight not match, random init')
# Print the model
print(netG)
# Create the encoder
netE = Encoder(args).to(device)
# Handle multi-gpu if desired
if (device.type == 'cuda') and (args.ngpu > 1):
netE = nn.DataParallel(netE, list(range(args.ngpu)))
# Apply the weights_init function to randomly initialize all weights
# to mean=0, stdev=0.2.
if args.model_id is 'default':
netE.apply(weights_init)
pass
else:
try:
pretrained_model = torch.load('./Encoders/' + args.model_id + '.pt')
try:
netE.load_state_dict(pretrained_model.state_dict())
# 1. 放入字典调用
def register(name=None, registry=None):
def decorator(fn, registration_name=None):
module_name = registration_name or _default_name(fn)
if module_name in registry:
raise LookupError(f"module {module_name} already registered.")
registry[module_name] = fn
return fn
return lambda fn: decorator(fn, name)
registry = {}
register = partial(register, registry=registry)
@register('simple')
class Encoder:
// 内部实现
from Encoders import registry as Encoder
encoder_name = config["encoder"]
encoder = Encoder[encoder_name](p1,p2,p3)
# 2. importlib
encoder_name = config["encoder"]
Encoder = importlib.import_module(encoder_name).component
encoder = Encoder(p1,p2,p3)