def __init__(self, s_size, a_size, random_seed):
"""Initialize an Agent object.
Params
======
s_size (int): dimension of each state (s)
a_size (int): dimension of each action (a)
random_seed (int): random seed
"""
self.s_size = s_size
self.a_size = a_size
self.seed = random.seed(random_seed)
# G: Generator (actor) Network (with Target Network)
self.g = G(s_size, a_size, random_seed).to(device)
self.g_target = G(s_size, a_size, random_seed).to(device)
self.g_optimizer = optim.Adam(self.g.parameters(), lr=LR)
# D: Discriminator (critic) or Decoder (predictor) Network (with Target Network)
self.d = D(s_size, a_size, random_seed).to(device)
self.d_target = D(s_size, a_size, random_seed).to(device)
self.d_optimizer = optim.Adam(self.d.parameters(), lr=LR)
# ReplayBuffer/ Memory
self.memory = Memory(a_size, BUFFER_SIZE, BATCH_SIZE, random_seed)
def generate():
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
netg, netd = G(DIM_NOISE, DIM_G).eval(), D(DIM_D).eval()
noises = torch.randn(GEN_SEARCH_NUM, DIM_NOISE, 1,
1).normal_(GEN_MEAN, GEN_STD)
noises = noises.to(device)
map_location = lambda storage, loc: storage
netd.load_state_dict(torch.load(netd_path, map_location=map_location))
netg.load_state_dict(torch.load(netg_path, map_location=map_location))
netd.to(device)
netg.to(device)
fake_img = netg(noises)
scores = netd(fake_img).detach()
# 挑选最好的某几张
indexs = scores.topk(GEN_NUM)[1]
result = []
for ii in indexs:
result.append(fake_img.data[ii])
# 保存图片
tv.utils.save_image(torch.stack(result),
"result.png",
normalize=True,
range=(-1, 1))
print('Successfully generate avatars')
def __init__(self, config: dict):
self.config = config
self.gen = G().to(device)
self.dis = D().to(device)
self.gen_op = torch.optim.Adam(self.gen.parameters(), lr=config['lr'])
self.dis_op = torch.optim.Adam(self.dis.parameters(), lr=config['lr'])
self.lda = config['lda']
self.epsilon = config['epsilon']
def build_model(self):
"""Builds a generator and a discriminator."""
from model import Dense_Net, Text_CNN_list, D
self.Gs = []
for i in range(self.n_view):
if i in self.text_views:
self.Gs.append(Text_CNN_list(mode=self.mode, word_dim=self.word_dim, vocab_size=self.vocab_size, out_dim=self.output_shape, filters=self.filters, filter_num=self.filter_num, dropout_prob=self.dropout_prob, wv_matrix=self.wv_matrix))
else:
self.Gs.append(Dense_Net(input_dim=self.input_shape[i], out_dim=self.output_shape))
self.D = D(dim=self.output_shape, view=self.n_view)
get_grad_params = lambda model: [x for x in model.parameters() if x.requires_grad]
g_params = [params for G in self.Gs for params in get_grad_params(G)]
d_params = get_grad_params(self.D)
self.g_optimizer = optim.Adam(g_params, self.lr, [self.beta1, self.beta2])
self.d_optimizer = optim.Adam(d_params, self.lr, [self.beta1, self.beta2])
if torch.cuda.is_available():
for G in self.Gs:
G.cuda()
self.D.cuda()
def train(iterations=10000,
batch_size=100,
sample_interval=5,
save_model_interval=100,
train_D_iters=1,
train_G_iters=3,
D_lr=0.0001,
G_lr=0.0001,
betas=(0.5, 0.99),
img_dir='./info_imgs',
model_dir='./models'):
imgs, digits, test_img, test_digits = load_mnist()
dataset = Dataset(imgs, digits)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
#dataset = Dataset(test_img, test_digits)
#test_loader = DataLoader(dataset,batch_size=batch_size, shuffle=False)
if torch.cuda.is_available:
print(f"Using GPU {torch.cuda.current_device()}")
device = "cuda"
else:
print("Using CPU...")
device = "cpu"
generaotor, discriminator, front, qq, encoder = G(), D(), FrontEnd(), Q(
), E()
generaotor = generaotor.to(device).apply(weights_init)
discriminator = discriminator.to(device).apply(weights_init)
qq = qq.to(device).apply(weights_init)
encoder = encoder.to(device).apply(weights_init)
front = front.to(device).apply(weights_init)
opt_G = torch.optim.Adam([{
"params": generaotor.parameters()
}, {
"params": qq.parameters()
}, {
"params": encoder.parameters()
}],
lr=G_lr,
betas=betas)
opt_D = torch.optim.Adam([{
"params": discriminator.parameters()
}, {
"params": front.parameters()
}],
lr=D_lr,
betas=betas)
CELoss_D = nn.CrossEntropyLoss(
weight=torch.FloatTensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])).to(device)
CELoss_G = nn.CrossEntropyLoss(weight=torch.FloatTensor(
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 20])).to(device)
CELoss_Q = nn.CrossEntropyLoss().to(device)
CosineLoss = nn.CosineEmbeddingLoss().to(device)
real_x = torch.FloatTensor(batch_size, 1, 32, 32).to(device)
trg = torch.LongTensor(batch_size).to(device)
label = torch.FloatTensor(batch_size, 1).to(device)
noise = torch.FloatTensor(batch_size, 54).to(device)
c = torch.FloatTensor(batch_size, 10).to(device)
v_target = torch.LongTensor(batch_size, 64).to(device) # For Q
real_x = Variable(real_x)
noise = Variable(noise)
c = Variable(c)
trg = Variable(trg, requires_grad=False)
label = Variable(label, requires_grad=False)
v_target = Variable(v_target, requires_grad=False)
for epoch in range(iterations):
for step, [batch_x, batch_target] in enumerate(loader):
bs = batch_x.size(0)
# train D
#==========
# real
opt_D.zero_grad()
real_x.data.copy_(batch_x)
trg.data.copy_(batch_target)
fe1 = front(real_x)
real_pred = discriminator(fe1)
real_loss = CELoss_D(real_pred, trg)
real_loss.backward()
#fake
real_x.data.copy_(batch_x)
v = encoder(real_x)
z, idx = noise_sample(c, noise, v, bs)
fake_stroke = generaotor(z)
fake_x = fake_stroke + real_x
fake_x = fake_x.clamp(max=1, min=0)
fe2 = front(fake_x.detach())
fake_pred = discriminator(fe2)
trg.data.fill_(10)
if epoch > 0:
ignore_rate = 0.01
else:
ignore_rate = 1
fake_loss = CELoss_D(fake_pred, trg) * ignore_rate
fake_loss.backward()
D_loss = real_loss + fake_loss
#D_loss.backward()
opt_D.step()
# train G, Q, E
#===============
#train G
opt_G.zero_grad()
fe = front(fake_x)
fake_pred = discriminator(fe)
trg.data.copy_(torch.LongTensor(idx))
reconstruct_loss = CELoss_G(fake_pred, trg)
# train Q
c_out, v_out = qq(fe)
class_ = torch.LongTensor(idx).to(device)
target = Variable(class_)
v_target.data.copy_(v)
# GQ Loss
q_c_loss = CELoss_Q(c_out, target)
q_v_loss = CosineLoss(v_out, v_target, label.data.fill_(1))
q_loss = q_c_loss + q_v_loss
G_loss = reconstruct_loss + q_c_loss + q_v_loss
G_loss.backward()
opt_G.step()
# accuracy
print(
f'Epoch: {epoch} | Dloss: {D_loss.data.cpu().numpy()} | QCloss: {q_c_loss.data.cpu().numpy()} | QVloss: {q_v_loss.data.cpu().numpy()} | reloss: {reconstruct_loss.data.cpu().numpy()}'
)
save_image(torch.cat((fake_x, fake_stroke, real_x), dim=0).data,
f'./{img_dir}/{epoch}.png',
nrow=20)
print(f"fake pred {np.argmax(fake_pred.data.cpu().numpy(),axis=1)}")
#print(f"Qpred {np.argmax(c_out[1].data.cpu().numpy())}")
#print(f"Origin {batch_target[1].data.cpu().numpy()} ToBe: {idx[0]}")
#save_image(real_x.data, f'./{img_dir}/{epoch}_R.png', nrow=10)
"""
train_data = DataLoader(dataset=train_set,
num_workers=opt.workers,
batch_size=opt.train_batch_size,
shuffle=True)
val_data = DataLoader(dataset=val_set,
num_workers=opt.workers,
batch_size=opt.test_batch_size,
shuffle=False)
print('=> Building model')
netG = G(opt.n_channel_input * 4, opt.n_channel_output,
opt.n_generator_filters)
netG.apply(weights_init)
netD = D(opt.n_channel_input * 4, opt.n_channel_output,
opt.n_discriminator_filters)
netD.apply(weights_init)
criterion = nn.BCELoss()
criterion_l1 = nn.L1Loss()
albedo = torch.FloatTensor(opt.train_batch_size, opt.n_channel_input, 256, 256)
direct = torch.FloatTensor(opt.train_batch_size, opt.n_channel_input, 256, 256)
normal = torch.FloatTensor(opt.train_batch_size, opt.n_channel_input, 256, 256)
depth = torch.FloatTensor(opt.train_batch_size, opt.n_channel_input, 256, 256)
gt = torch.FloatTensor(opt.train_batch_size, opt.n_channel_output, 256, 256)
label = torch.FloatTensor(opt.train_batch_size)
real_label = 1
fake_label = 0
def __init__(self, opt, root_dir):
self.opt = opt
print('=> Loading datasets')
train_dir = join(root_dir + self.opt.dataset, "train")
test_dir = join(root_dir + self.opt.dataset, "val")
train_set = DataLoaderHelper(train_dir)
val_set = DataLoaderHelper(test_dir)
self.batch_size = self.opt.train_batch_size
self.n_epoch = self.opt.n_epoch
self.lastEpoch = 0
# workers == 0, run main process, default 0
self.train_data = DataLoader(dataset=train_set,
num_workers=self.opt.workers,
batch_size=self.opt.train_batch_size,
shuffle=True)
self.val_data = DataLoader(dataset=val_set,
num_workers=self.opt.workers,
batch_size=self.opt.test_batch_size,
shuffle=False)
# debug
"""
for (i, images) in enumerate(self.train_data):
(albedo_cpu, direct_cpu, normal_cpu, depth_cpu, gt_cpu) = (images[0], images[1], images[2], images[3], images[4])
# debug (1,3,256,256)
debug_albedo = albedo_cpu[0]
debug_albedo = debug_albedo.add_(1).div_(2)
debug_albedo = debug_albedo.numpy()
print(debug_albedo.shape)
debug_albedo *= 255.0
debug_albedo = debug_albedo.clip(0, 255)
debug_albedo = np.transpose(debug_albedo, (1, 2, 0)) # (width, height, channels)
debug_albedo = debug_albedo.astype(np.uint8)
plt.imshow(debug_albedo)
plt.show()
print(debug_albedo.shape)
"""
print('=> Building model')
self.netG = G(self.opt.n_channel_input * 4, self.opt.n_channel_output,
self.opt.n_generator_filters)
self.netG.apply(weights_init)
self.netD = D(self.opt.n_channel_input * 4, self.opt.n_channel_output,
self.opt.n_discriminator_filters)
self.netD.apply(weights_init)
self.criterion = nn.BCELoss()
self.criterion_l1 = nn.L1Loss()
self.label = torch.FloatTensor(self.opt.train_batch_size)
self.real_label = 1
self.fake_label = 0
self.albedo = torch.FloatTensor(self.opt.train_batch_size,
self.opt.n_channel_input, 256, 256)
self.direct = torch.FloatTensor(self.opt.train_batch_size,
self.opt.n_channel_input, 256, 256)
self.normal = torch.FloatTensor(self.opt.train_batch_size,
self.opt.n_channel_input, 256, 256)
self.depth = torch.FloatTensor(self.opt.train_batch_size,
self.opt.n_channel_input, 256, 256)
self.gt = torch.FloatTensor(self.opt.train_batch_size,
self.opt.n_channel_output, 256, 256)
# GPU
self.netD = self.netD.cuda()
self.netG = self.netG.cuda()
self.criterion = self.criterion.cuda()
self.criterion_l1 = self.criterion_l1.cuda()
self.albedo = self.albedo.cuda()
self.direct = self.direct.cuda()
self.normal = self.normal.cuda()
self.depth = self.depth.cuda()
self.gt = self.gt.cuda()
self.label = self.label.cuda()
# Derivative
self.albedo = Variable(self.albedo)
self.direct = Variable(self.direct)
self.normal = Variable(self.normal)
self.depth = Variable(self.depth)
self.gt = Variable(self.gt)
self.label = Variable(self.label)
# Optimizer
self.optimizerD = optim.Adam(self.netD.parameters(),
lr=self.opt.lr,
betas=(self.opt.beta1, 0.999))
self.optimizerG = optim.Adam(self.netG.parameters(),
lr=self.opt.lr,
betas=(self.opt.beta1, 0.999))
if self.opt.resume_G:
if os.path.isfile(self.opt.resume_G):
print("=> loading generator checkpoint '{}'".format(
self.opt.resume_G))
checkpoint = torch.load(self.opt.resume_G)
self.lastEpoch = checkpoint['epoch']
self.n_epoch = self.n_epoch - self.lastEpoch
self.netG.load_state_dict(checkpoint['state_dict_G'])
self.optimizerG.load_state_dict(checkpoint['optimizer_G'])
print("=> loaded generator checkpoint '{}' (epoch {})".format(
self.opt.resume_G, checkpoint['epoch']))
else:
print("=> no checkpoint found")
if self.opt.resume_D:
if os.path.isfile(self.opt.resume_D):
print("=> loading discriminator checkpoint '{}'".format(
self.opt.resume_D))
checkpoint = torch.load(self.opt.resume_D)
self.netD.load_state_dict(checkpoint['state_dict_D'])
self.optimizerD.load_state_dict(checkpoint['optimizer_D'])
print("=> loaded discriminator checkpoint '{}'".format(
self.opt.resume_D))
imgSize = 32
batchSize = 16
criterion = nn.BCELoss()
device = torch.device("cuda" if cuda else "cpu")
path = '/run/media/why/DATA/why的程序测试/AI_Lab/Task/task_week6/img'
trainData = dSet.MNIST(root='/run/media/why/DATA/why的程序测试/AI_Lab/DataSet',
train=True,
transform=T.Compose([T.Resize(imgSize),
T.ToTensor()]),
download=True) # size = 28*28
trainLoader = torch.utils.data.DataLoader(dataset=trainData,
batch_size=batchSize,
shuffle=True)
gNet = G()
dNet = D()
if cuda:
dNet = dNet.cuda()
gNet = gNet.cuda()
D_optim = optim.Adam(dNet.parameters(), lr=lr, betas=(0.5, 0.999))
G_optim = optim.Adam(gNet.parameters(), lr=lr, betas=(0.5, 0.999))
fixed_noise = torch.randn(batchSize, nz, 1, 1).cuda() # 用来生成假图片的噪声
for j in range(epoch):
for i, (img, _) in enumerate(trainLoader, 0):
# 训练识别器：
dNet.zero_grad() # 初始化所有梯度
real = img.to(device) # 训练集中抽取的真图片
label = torch.full((batchSize, ), 1, device=device) # 为所有真图片标记1
output = dNet(real)
err_D = criterion(output, label) # 计算loss
dis_c.data.copy_(torch.Tensor(one_hot))
print(np.shape(c1))
con_c = Variable(torch.rand(con_c.size())).cuda()
z = torch.cat([noise, dis_c, con_c], 1).view(-1, 74, 1, 1)
x_save = self.G(z)
save_image(x_save.data, os.path.join(args.path, 'generate.png'), nrow=10)
def parse():
parser = argparse.ArgumentParser(description='VAE MNIST Example')
parser.add_argument('--label', type=int, default=1, metavar='N',
help='The label you want to generate')
parser.add_argument('--num', type=int, default=1, metavar='N',
help='The number of image you want to generate')
args = parser.parse_args()
args.cuda = torch.cuda.is_available()
args.path = './infoGAN_result'
return args
if __name__ == '__main__':
args = parse()
fe = FrontEnd()
d = D()
q = Q()
g = G()
for i in [fe, d, q, g]:
i.cuda()
tester = Tester(g, fe, d, q, args)
tester.load()
tester.generate()
batch_size=..., shuffle=True)
val_data = DataLoader(dataset=val_set, num_workers=...,
batch_size=..., shuffle=False)
test_data = DataLoader(dataset=test_set,...)
# cast loaded VTK-m data to PyTorch tensors
gt = torch.FloatTensor(train_batch_size,
n_channel_input,
256, 256)
direct, albedo, depth, normal = torch.FloatTensor(...)
, ... , ... , ...
# ground truth labels
label = torch.FloatTensor(train_batch_size)
# instantiate generator and descriminator
netG = G(n_channel_input*4, n_channel_output,
n_generator_filters)
netD = D(n_channel_input*4, n_channel_output,
n_discriminator_filters)
# assign to GPU
netD, netG = netD.cuda(), netG.cuda()
# loss functions:
# Binary Cross Entropy and L1 Loss
criterion, criterion_l1 = nn.BCELoss(), nn.L1Loss()
# Cuda placement
criterion = criterion.cuda()
criterion_l1 = criterion_l1.cuda()
albedo = albedo.cuda()
gt, direct = gt.cuda(), direct.cuda()
depth, normal = depth.cuda(), normal.cuda()
label = label.cuda()
# instantiate PyTorch variables of VTK-m renderings
albedo = Variable(albedo)
gt, direct = Variable(gt), Variable(direct)
c_2[range(bs), idx_2] = 1.0
# print('c_2: ', c_2)
dis_c.data.copy_(torch.Tensor(c))
con_c.data.copy_(torch.Tensor(c_2))
noise.data.uniform_(-1.0, 1.0)
print('noise: ', noise.shape)
z = torch.cat([noise, dis_c, con_c], 1).view(-1, 74, 1, 1)
return z, idx, idx_2
model_Q = Q().to(device)
model_FE = FrontEnd().to(device)
model_G = G().to(device)
model_D = D().to(device)
model_Q.load_state_dict(
torch.load(model_path + '/model_Q.pytorch', map_location='cpu'))
model_D.load_state_dict(
torch.load(model_path + '/model_D.pytorch', map_location='cpu'))
model_FE.load_state_dict(
torch.load(model_path + '/model_FE.pytorch', map_location='cpu'))
model_G.load_state_dict(
torch.load(model_path + '/model_G.pytorch', map_location='cpu'))
model_Q.eval()
model_D.eval()
model_FE.eval()
model_G.eval()
def train(config):
################################### Model Initalization ####################################
dataset = Dataset(config['dir_csv']
) # Class declaraction --> Runs ___init___ with Dataset
dataloader = DataLoader(dataset,
batch_size=1) # Built-in class within Pytorch
'''
Save Ground Truth Plots
'''
if True:
Path('ground_truths').mkdir(parents=True, exist_ok=True)
for data in dataloader:
real_data = data['real_data']
basename = data['basename'][0]
save_intermediate_results(
os.path.join('ground_truths', basename + '.jpg'), real_data)
gen = G().to(config['device'])
dsc = D().to(config['device'])
optimizer_G = torch.optim.Adam(gen.parameters(), lr=config['lr'])
optimizer_D = torch.optim.Adam(dsc.parameters(), lr=config['lr'])
real_label = torch.tensor(1.0).view(1, -1).to(
config['device']) # Tensor of shape (1, 1)
fake_label = torch.tensor(0.0).view(1, -1).to(
config['device']) # Tensor of shape (1, 1)
criterion = nn.BCELoss() # Binary Cross Entropy Loss
fixed_noise = torch.rand((1, 100)).to(config['device'])
for epoch in range(config['n_epoch']):
for data in dataloader:
real_data = data['real_data'].to(config['device'])
##################### Optimize for Generator ##########################
optimizer_G.zero_grad()
fake_data = gen(fixed_noise) # (1, 100) -> (1, 1, 800)
pred = dsc(fake_data) # (1, 1, 800) -> (1, 1)
G_loss = criterion(
pred,
real_label) # Train the generator to fool the discriminator
'''
Optimize
'''
G_loss.backward()
optimizer_G.step()
##################### Optimize for Discriminator ######################
optimizer_D.zero_grad()
'''
Real Input
'''
pred = dsc(real_data) # (1, 1, 800) -> (1, 1)
D_loss_real = criterion(
pred, real_label
) # Train the discriminator to distinguish between real and fake data
'''
Fake Input
'''
pred = dsc(fake_data.detach()) # (1, 1, 800) -> (1, 1)
D_loss_fake = criterion(
pred, fake_label
) # Train the discriminator to distinguish between real and fake data
'''
Optimize
'''
D_loss_total = (D_loss_real + D_loss_fake) / 2
D_loss_total.backward()
optimizer_D.step()
if (((epoch + 1) % config['val_epoch']) == 0):
Path('results').mkdir(parents=True, exist_ok=True)
save_intermediate_results(
os.path.join('results', 'epoch_%d.jpg' % (epoch + 1)),
fake_data.detach().cpu())
print('[Epoch] %d / %d' % (epoch + 1, config['n_epoch']), end='\r')
def train():
device = torch.device('cuda') if torch.cuda.is_available() else 'cpu'
transforms = tv.transforms.Compose([
tv.transforms.Resize(IMG_SIZE),
tv.transforms.CenterCrop(IMG_SIZE),
tv.transforms.ToTensor(),
tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
vis = Visualizer(env)
# drop_last: retain data that cannot suffice a batch
dataset = tv.datasets.ImageFolder(DATA_PATH, transforms)
dataloader = data.DataLoader(dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=4,
drop_last=True)
netG, netD = G(DIM_NOISE, DIM_G), D(DIM_D)
netG.to(device)
netD.to(device)
# Optimizers and Loss functions
optimizer_G = torch.optim.Adam(netG.parameters(),
lr=LR_G,
betas=(0.5, 0.999))
optimizer_D = torch.optim.Adam(netG.parameters(),
lr=LR_D,
betas=(0.5, 0.999))
criterion = nn.BCELoss().to(device)
noises = torch.randn(BATCH_SIZE, DIM_NOISE, 1, 1).to(device)
errord_meter = AverageValueMeter()
errorg_meter = AverageValueMeter()
# Real imgs have 0 loss; fake imgs have 1
true_labels = torch.ones(BATCH_SIZE).to(device)
fake_labels = torch.zeros(BATCH_SIZE).to(device)
max_epoch_G, max_epoch_D = 0, 0
# Load pretrained network with greatest number of training epochs
pretrained_root_g = MODEL_SAVE_PATH + 'G/'
if PRETRAINED and os.path.exists(pretrained_root_g):
file_list = [
file for file in os.listdir(pretrained_root_g)
if file.endswith('.pth')
]
if file_list != []:
index_list = [
int(file.split('.pth')[0].split('_')[2]) for file in file_list
]
max_epoch_G = max(index_list)
model_name = 'model_g_%s.pth' % max_epoch_G
print('Using mode:', model_name)
netG.load_state_dict(torch.load(pretrained_root_g + model_name))
netG.to(device)
else:
print('Generator train from Step 0')
pretrained_root_d = MODEL_SAVE_PATH + 'D/'
if PRETRAINED and os.path.exists(pretrained_root_d):
file_list = [
file for file in os.listdir(pretrained_root_d)
if file.endswith('.pth')
]
if file_list != []:
index_list = [
int(file.split('.pth')[0].split('_')[2]) for file in file_list
]
max_epoch_D = max(index_list)
model_name = 'model_d_%s.pth' % max_epoch_D
print('Using mode:', model_name)
netD.load_state_dict(torch.load(pretrained_root_d + model_name))
netD.to(device)
else:
print('Discriminator train from Epoch 0')
for epoch in range(EPOCHS):
time_start = time.time()
for step, (img, _) in enumerate(dataloader):
# Skip former steps if using pretrained models
real_img = img.to(device)
if step <= max_epoch_G and step <= max_epoch_D:
continue
# Train generator
if step % GENERATE_EVERY == 0 and step > max_epoch_G:
print('G - Epoch:', epoch, '| Step:', step)
optimizer_G.zero_grad()
noises.data.copy_(
torch.randn(BATCH_SIZE, DIM_NOISE, 1, 1)
) # TODO: Why not noises=(torch.randn(BATCH_SIZE, DIM_NOISE, 1, 1))
fake_img = netG(noises)
output_g = netD(fake_img)
loss_g = criterion(output_g, true_labels)
loss_g.backward()
optimizer_G.step()
errorg_meter.add(loss_g.item())
if step % DISCRIMINATE_EVERY == 0 and step > max_epoch_D:
# Identify real images
print('D - Epoch:', epoch, '| Step:', step)
output_d = netD(real_img)
loss_real = criterion(output_d, true_labels)
loss_real.backward()
# Identify fake images
optimizer_D.zero_grad()
noises.data.copy_(torch.randn(BATCH_SIZE, DIM_NOISE, 1, 1))
fake_img = netG(noises)
output_d = netD(fake_img)
loss_fake = criterion(output_d, fake_labels)
loss_fake.backward()
optimizer_D.step()
loss_d = loss_real + loss_fake
errord_meter.add(loss_d.item())
#fake_img = fake_img.detach()
if (step + 1) % PLOT_EVERY:
fix_fake_imgs = netG(noises)
vis.images(fix_fake_imgs.detach().cpu().numpy()[:64] * 0.5 +
0.5,
win='fixfake')
vis.images(real_img.data.cpu().numpy()[:64] * 0.5 + 0.5,
win='real')
vis.plot('errord', errord_meter.value()[0])
vis.plot('errorg', errorg_meter.value()[0])
time_end = time.time()
print('Total time for epoch ', epoch, ' is ', (time_end - time_start))
if epoch and epoch % SAVE_EVERY == 0:
torch.save(netG.state_dict(),
MODEL_SAVE_PATH + 'G/model_g_%s.pth' % step)
torch.save(netD.state_dict(),
MODEL_SAVE_PATH + 'D/model_d_%s.pth' % step)
errord_meter.reset()
errorg_meter.reset()
def build_model(gpu, gpu_per_node, args):
args.gpu = gpu
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
log.write("Use GPU: {} for training".format(args.gpu))
if not opt.multiprocess_distributed:
if torch.cuda.device_count() > 1:
opt.world_size = torch.cuda.device_count() #
#args.gpu = gpu_per_node
#else:
# if torch.cuda.device_count() > 1:
# gpu = None
# args.gpu = gpu_per_node
device_G = None
device_D = None
device_G_id = None
device_D_id = None
if args.multiprocess_distributed:
args.rank = args.rank * gpu_per_node + gpu
dist.init_process_group(backend=args.dist_backend,world_size=args.world_size, rank=args.rank)
print('=> Building model')
if not opt.multiprocess_distributed:
netG = G(opt.n_channel_input*4, opt.n_channel_output, opt.n_generator_filters)
netD = D(opt.n_channel_input*4, opt.n_channel_output, opt.n_discriminator_filters)
# specify gpus to use for data parallelism
# setup devices for this process, split 1 uses GPUs [0, 1, 2, 3] and
# rank 2 uses GPUs [4, 5, 6, 7].
def split_gpu(devices, split):
n = devices // args.world_size
device_ids = list(range(split * n, (split + 1) * n))
return device_ids
device_G, device_D = [] , []
if args.multiprocess_distributed:
print("using multi process distributed training")
print("taining with GPU: ", args.gpu)
torch.cuda.set_device(args.gpu)
#netG.cuda(args.gpu)
#netD.cuda(args.gpu)#cuda(device_D)
device_G_id = args.gpu
#
# strongly distributed
#args.train_batch_size = int(args.train_batch_size / gpu_per_node)
#args.test_batch_size = int(args.test_batch_size / gpu_per_node)
#args.workers = int((args.workers + gpu_per_node - 1) / gpu_per_node)
#
args.workers = 0
print("Let's use", torch.cuda.device_count(), "GPUs!")
node_gpus = list(range( torch.cuda.device_count()))
gan_process = torch.distributed.new_group(range(dist.get_world_size()))
netG = torch.nn.parallel.DistributedDataParallel(G(opt.n_channel_input*4
, opt.n_channel_output
, opt.n_generator_filters)
,device_ids = [args.gpu]
,process_group = gan_process)
dis_process = torch.distributed.new_group(range(dist.get_world_size()))
netD = torch.nn.parallel.DistributedDataParallel(D(opt.n_channel_input*4
, opt.n_channel_output
, opt.n_discriminator_filters).cuda(args.gpu)
,device_ids = [args.gpu]
,process_group = dis_process)
netG.cuda(args.gpu)
netD.cuda(args.gpu)
if gpu is not None:
print("Setting to ", gpu, " GPU.")
#torch.cuda.set_device(gpu)
#model = model.cuda(args.gpu)
device_G_id = gpu[0]
args.gpu = gpu[0]
netG = nn.DataParallel(netG).to(gpu[0])
netD = nn.DataParallel(netD).to(gpu[0])
else:
if torch.cuda.device_count() > 1:
opt.world_size = torch.cuda.device_count()
device_G = split_gpu(torch.cuda.device_count(), 1)
device_D = split_gpu(torch.cuda.device_count(), 2)
print("Using non-distributed parallelized over device ",torch.cuda.device_count(), " GPU.")
print("generative net on gpu ", device_G[0])
print("descriminator on gpu ", device_D[0])
device_G_id = device_G[0]
device_D_id = device_D[0]
#args.gpu = device_G[0]
#netG.to(device_G[0])#cuda(device_G)
#netD.to(device_G[0])
#netG.features = torch.nn.DataParallel(netG.features)
#netD.features = torch.nn.DataParallel(netD.features)
netG = nn.DataParallel(netG, device_G+device_D).cuda(device_G[0])#.cuda()
netD = nn.DataParallel(netD, device_G+device_D).cuda(device_G[0])#.cuda()
else:
print("Using device ",torch.cuda.device_count(), " GPU.")
device_G_id = device
args.gpu = device
netG = nn.DataParallel(netG).to(device)
netD = nn.DataParallel(netD).to(device)
#netD.to(device)
#netG.to(device)
netG.apply(dynamic_weights_init)
netD.apply(weights_init)
criterion = nn.BCELoss()#.cuda(device_G_id)#args.gpu)
criterion_l1 = nn.L1Loss()#.cuda(device_G_id)#args.gpu) # to use multiple gpu among multipl model
#cuda auto-tuner to find best algorithm given hardware
cudnn.benchmark = True
print('=> Loading Data')
root_dir = opt.dataset#"../PathTracer/build"
conditional_names = ["outputs", "direct", "depth", "normals", "albedo"]
# some higher resolution images
proj_read_dir = os.path.join(os.environ["PROJWORK"],"csc143","samlev","data")
train_set = AdiosDataLoader( os.path.join(proj_read_dir, opt.dataset),split="train") #DataLoaderHelper(train_dir)
val_set = AdiosDataLoader( os.path.join(proj_read_dir, opt.dataset), split="val") #DataLoaderHelper(test_dir)
test_set = AdiosDataLoader( os.path.join(proj_read_dir, opt.dataset), split="test")
batch_size = opt.train_batch_size
n_epoch = opt.n_epoch
args.workers = 0
pin_memory = True
# Change if used in paper!
if opt.multiprocess_distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_set)
pin_memory = True
#val_sampler = torch.utils.data.distributed.DistributedSampler(val_set)
#test_sampler = torch.utils.data.distributed.DistributedSampler(test_set)
else:
train_sampler = None
pin_memory = False
train_data = DataLoader(dataset=train_set
, batch_size=args.train_batch_size
,shuffle=(train_sampler is None)
,num_workers=args.workers
, pin_memory=pin_memory
, sampler=train_sampler)
val_data = DataLoader(dataset=val_set
,num_workers = args.workers
, batch_size=opt.test_batch_size
, shuffle=False
,pin_memory=pin_memory)
test_data = DataLoader(dataset=test_set
, num_workers = args.workers
, batch_size=opt.test_batch_size
, shuffle=False
,pin_memory=pin_memory)
albedo = torch.FloatTensor(args.train_batch_size, opt.n_channel_input, 256, 256)
direct = torch.FloatTensor(args.train_batch_size, opt.n_channel_input, 256, 256)
normal = torch.FloatTensor(args.train_batch_size, opt.n_channel_input, 256, 256)
depth = torch.FloatTensor(args.train_batch_size, opt.n_channel_input, 256, 256)
gt = torch.FloatTensor(args.train_batch_size, opt.n_channel_output, 256, 256)
label = torch.FloatTensor(args.train_batch_size)
fake_label = 0.1#numpy.random.uniform(0. , 0.1)
real_label = 0.9# numpy.random.uniform(0.9, 1.0)
#for multigpu multimodel assign gpu to whole model
# with model already assigned to gpu devices.
# data parrallel later
if (torch.cuda.device_count() == 1):# or ((not opt.multiprocess_distributed) and (torch.cuda.device_count() > 1)):
print("assigning non threaded data")
albedo = albedo.to(device_G_id)#.cuda()
direct = direct.to(device_G_id)#.cuda()
normal = normal.to(device_G_id)#.cuda()
depth = depth.to(device_G_id)#.cuda()
gt = gt.to(device_G_id)#.cuda()
label = label.to(device_G_id)#.cuda()
#albedo = Variable(albedo)
#direct = Variable(direct)
#normal = Variable(normal)
#depth = Variable(depth)
#gt = Variable(gt)
#label = Variable(label)
elif torch.cuda.device_count() > 1:
print("assigning distributed data")
albedo = albedo.to(device_G_id, non_blocking=True)#to(device)#.cuda()
direct = direct.to(device_G_id, non_blocking=True)#.to(device)#.cuda()
normal = normal.to(device_G_id, non_blocking=True)#.to(device)#.cuda()
depth = depth.to(device_G_id, non_blocking=True)#.to(device)#.cuda()
gt = gt.to(device_G_id, non_blocking=True)#.to(device)#.cuda()
label = label.to(device_G_id, non_blocking=True)#.to(device)#.cuda()
albedo = Variable(albedo)
direct = Variable(direct)
normal = Variable(normal)
depth = Variable(depth)
gt = Variable(gt)
label = Variable(label)
optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
if opt.resume_G:
if os.path.isfile(opt.resume_G):
print("=> loading generator checkpoint '{}'".format(opt.resume_G))
checkpoint = torch.load(opt.resume_G)
lastEpoch = checkpoint['epoch']
n_epoch = n_epoch - lastEpoch
netG.load_state_dict(checkpoint['state_dict_G'])
optimizerG.load_state_dict(checkpoint['optimizer_G'])
print("=> loaded generator checkpoint '{}' (epoch {})".format(opt.resume_G, checkpoint['epoch']))
else:
print("=> no checkpoint found")
if opt.resume_D:
if os.path.isfile(opt.resume_D):
print("=> loading discriminator checkpoint '{}'".format(opt.resume_D))
checkpoint = torch.load(opt.resume_D)
netD.load_state_dict(checkpoint['state_dict_D'])
optimizerD.load_state_dict(checkpoint['optimizer_D'])
print("=> loaded discriminator checkpoint '{}'".format(opt.resume_D))
#time profiling
# time profiling
# profiling and logging
filename = 'time_profiling_'+run_name+'.txt'
filename = os.path.join(proj_write_dir, filename)
if os.path.exists(filename):
append_write = 'a' # append if already exists
else:
append_write = 'w' # make a new file if not
tlog = open(filename,append_write)
time0 = time.time()
tlog.write(str(time0))
tlog.close()
lastEpoch = 0
# begin training
for epoch in range(opt.n_epoch):
train(epoch
, train_data
, albedo, direct, normal, depth, gt, label
, netG, netD
, criterion, criterion_l1
, optimizerG, optimizerD)
if epoch % 5 == 0:
save_checkpoint(epoch+lastEpoch
, train_data
, val_data
, albedo, direct, normal, depth, gt, label
, netG, netD
, criterion, criterion_l1
, optimizerG, optimizerD)
tlog = open(filename,append_write)
tlog.write(str(time0 - time.time()))
tlog.close()
test_set = get_test_set(root_path + opt.dataset)
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True)
testing_data_loader = DataLoader(dataset=test_set,
num_workers=opt.threads,
batch_size=opt.testBatchSize,
shuffle=False)
# ======= Load the Generator and Discriminator ===============
print('===> Building model')
netG = G(opt.input_nc, opt.output_nc, opt.ngf)
pytorch_total_params = sum(p.numel() for p in netG.parameters()
if p.requires_grad)
print("\nTrainable parameters", pytorch_total_params)
netD = D(opt.input_nc, opt.output_nc, opt.ndf)
#========Loss Function==========
criterion = nn.BCELoss()
criterion_l1 = nn.L1Loss()
criterion_mse = nn.MSELoss()
#================================
real_A = torch.FloatTensor(opt.batchSize, opt.input_nc, 128, 128)
real_B = torch.FloatTensor(opt.batchSize, opt.output_nc, 128, 128)
label = torch.FloatTensor(opt.batchSize)
real_label = 1
fake_label = 0
if opt.cuda:
netD = netD.cuda()
netG = netG.cuda()
criterion = criterion.cuda()
def train(config):
face_M_dataset = CusDataset(csv_file=config.csv_dir,
root_dir=config.root_dir,
Wem=0,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))
]))
face_W_dataset = CusDataset(csv_file=config.csv_dir,
root_dir=config.root_dir,
Wem=1,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))
]))
M_loader = DataLoader(face_M_dataset,
batch_size=config.batch_size,
shuffle=True,
num_workers=4)
W_loader = DataLoader(face_W_dataset,
batch_size=config.batch_size,
shuffle=True,
num_workers=4)
Dw = D()
Dm = D()
Gwm = G()
Gmw = G()
if torch.cuda.is_available():
print('cuda')
Dw = Dw.cuda()
Dm = Dm.cuda()
Gwm = Gwm.cuda()
Gmw = Gmw.cuda()
g_params = list(Gwm.parameters()) + list(Gmw.parameters())
d_params = list(Dw.parameters()) + list(Dm.parameters())
#print(d_params)
g_optimer = optim.Adam(g_params, lr=config.lr)
d_optimer = optim.Adam(d_params, lr=config.lr)
criterion = torch.nn.CrossEntropyLoss()
writer = SummaryWriter()
global_step = 0
for i in range(config.epoch):
M_data = iter(M_loader)
W_data = iter(W_loader)
while True:
try:
M_batched = M_data.__next__()
W_batched = W_data.__next__()
except StopIteration:
break
M_image = M_batched['imagename'].type(torch.cuda.FloatTensor)
M_label = M_batched['age'].type(torch.cuda.LongTensor).squeeze()
W_image = W_batched['imagename'].type(torch.cuda.FloatTensor)
W_label = W_batched['age'].type(torch.cuda.LongTensor).squeeze()
if W_label.size(0) != M_label.size(0):
break
fake_label = np.zeros([W_label.size(0)])
fake_label[:] = 10
#fake_label[:,10] = 1
fake_label = torch.from_numpy(fake_label)
fake_label = fake_label.type(torch.cuda.LongTensor)
#print(M_label.type())
'''print(fake_label.size())
print("\n{}".format(W_label.view(W_label.size(0)*10)))
'''
#-----------------------------train D---------------------------
#-------------train real----------
g_optimer.zero_grad()
d_optimer.zero_grad()
dm_real_loss = criterion(Dm(M_image), M_label)
dw_real_loss = criterion(Dw(W_image), W_label)
d_real_loss = dm_real_loss + dw_real_loss
d_real_loss.backward()
d_optimer.step()
#-------------train fake-----------
g_optimer.zero_grad()
d_optimer.zero_grad()
dm_fake_loss = criterion(Dm(Gwm(W_image)), fake_label)
dw_fake_loss = criterion(Dw(Gmw(M_image)), fake_label)
d_fake_loss = dm_fake_loss + dw_fake_loss
d_fake_loss.backward()
d_optimer.step()
#-----------------------------train G-----------------------------
#----------------train--W-M-W-Cycle-------------
g_optimer.zero_grad()
d_optimer.zero_grad()
fake_m_image = Gwm(W_image)
g_WMW_loss = criterion(Dm(fake_m_image), M_label) + torch.mean(
((W_image - Gmw(fake_m_image))**2))
g_WMW_loss.backward()
g_optimer.step()
#----------------train--M-W-M-Cycle-------------
g_optimer.zero_grad()
d_optimer.zero_grad()
fake_w_image = Gmw(M_image)
g_MWM_loss = criterion(Dw(fake_w_image), W_label) + torch.mean(
((M_image - Gwm(fake_w_image))**2))
g_MWM_loss.backward()
g_optimer.step()
fake_w_image = vutils.make_grid(fake_w_image,
normalize=False,
scale_each=False)
fake_m_image = vutils.make_grid(fake_m_image,
normalize=False,
scale_each=False)
M_image = vutils.make_grid(M_image,
normalize=False,
scale_each=False)
writer.add_scalar(tag='loss/g_MWM_loss',
scalar_value=g_MWM_loss,
global_step=global_step)
writer.add_scalar(tag='loss/g_WMW_loss',
scalar_value=g_WMW_loss,
global_step=global_step)
writer.add_scalar(tag='loss/d_fake_loss',
scalar_value=d_fake_loss,
global_step=global_step)
writer.add_scalar(tag='loss/d_real_loss',
scalar_value=d_real_loss,
global_step=global_step)
writer.add_image('image/fake_w_image', fake_w_image, global_step)
writer.add_image('image/fake_m_image', fake_m_image, global_step)
writer.add_image('image/M_image', M_image, global_step)
print("{}\n".format(global_step))
global_step += 1
writer.close()
'''for i_batch,sample_batched in enumerate(dataloader):
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
train_dataset = dataset.MNIST(root='./data/',
train=True,
transform=transform,
download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
gx = P()
gz = Q()
dxz = D()
g_param = chain(gx.parameters(), gz.parameters())
d_param = dxz.parameters()
g_optimizer = optim.Adam(g_param, glr, betas=(0.5, 0.999))
d_optimizer = optim.Adam(d_param, dlr, betas=(0.5, 0.999))
def to_variable(x):
if torch.cuda.is_available():
x = x.cuda()
return Variable(x)
def to_data(x):