if __name__ == "__main__":
# devie---------------------------------------------------------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Fix random seed---------------------------------------------------------------------------
# torch.manual_seed(2)
# torch.cuda.manual_seed_all(2)
ber_list = []
for i in range(1):
# model------------------------------------------------------------------------------------
stylegan = build_model(args.experiment,
image_size=args.image_size,
lr=args.lr)
stylegan = checkpointNet.load_part_network(stylegan,
args.checkpoint_GAN,
args.which_epoch)
stylegan = stylegan.to(device)
# criterion-----------------------------------------------------------------------------------
criterion = nn.MSELoss()
# optimizer-----------------------------------------------------------------------------------
for p in stylegan.NE.parameters():
p.requires_grad = True
param_groups = [{'params': stylegan.NE.parameters(), 'lr': args.lr}]
optimizer = torch.optim.Adam(param_groups,
lr=args.lr,
betas=(0.5, 0.999))
# scheduler-----------------------------------------------------------------------------------
import sys
sys.path.append('/home/hy/vscode/StyleGANStege')
import torch
from models import build_model
from utils import checkpointNet
if __name__ == "__main__":
checkpoint = '/home/hy/vscode/StyleGANStege/experiments/ExtractNet'
model = build_model('ExtractNet', image_size=64)
model = checkpointNet.load_part_network(model, checkpoint, 'final')
# print(model)
mixed_latents = [(tmp1, tt), (tmp2, num_layers - tt)]
# generated_images--------------------------------------------
generated_images = generate_truncated(model.SE, model.GE, mixed_latents, n,
av, batch_size)
fo.write(str(noise_.numpy()[0]))
torchvision.utils.save_image(
generated_images,
str(Path(save_dir_path) / f'{str(num)}-mr.{ext}'),
nrow=5)
if __name__ == "__main__":
# Fix random seed---------------------------------------------------------------------------
torch.manual_seed(1)
torch.cuda.manual_seed_all(1)
# model------------------------------------------------------------------------------------
checkpoint = '/home/hy/vscode/StyleGANStege/experiments/Celeba'
which_epoch = 'final'
experiment = 'StyleGAN2'
model = build_model(experiment, image_size=64, lr=1)
model = checkpointNet.load_part_network(model, checkpoint, which_epoch)
model = model.to(device)
# save_dir_path-----------------------------------------------------------------------------------
save_path = './experiments'
save_dir_path = os.path.join(save_path, 'TrainEImg')
os.makedirs(save_dir_path, exist_ok=True)
genimg(model, save_dir_path)
type=str,
default='/home/hy/vscode/StyleGANStege/experiments/Celeba')
parser.add_argument('--dataset', type=str, default='Celeba')
parser.add_argument('--batch_size', default=3, type=int, help='batch_size')
parser.add_argument('--image_size', default=64)
parser.add_argument('--lr', type=float, default=0.1)
args = parser.parse_args()
# devie---------------------------------------------------------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# save_dir_path-----------------------------------------------------------------------------------
save_dir_path = os.path.join(args.save_path, args.dataset)
os.makedirs(save_dir_path, exist_ok=True)
print('save path is : {}'.format(save_dir_path))
# model------------------------------------------------------------------------------------
model = build_model(args.experiment,
image_size=args.image_size,
lr=args.lr)
model = checkpointNet.load_part_network(model, args.checkpoint,
args.which_epoch)
model = model.to(device)
test(model, save_dir_path, args)
# torch.cuda.empty_cache()