def getPsnr(fake, real, use_gpu=False):
mse_loss = MSELoss()
if use_gpu:
mse_loss = mse_loss.cuda()
mse = mse_loss(fake, real)
psnr = 10 * log10(1 / mse.data.item())
return psnr
def main(args):
if args.model == 'one':
Net = One
if args.model == 'two':
Net = Two
if args.model == 'unet':
Net = UNet
model = Net()
crit1 = L1Loss()
crit2 = MSELoss()
if args.cuda:
model = model.cuda()
crit1 = crit1.cuda()
crit2 = crit2.cuda()
if args.state:
model.load_state_dict(torch.load(args.state))
trainer = Trainer(args, model)
optimizer = Adam(model.parameters())
test_loader = DataLoader(
MNIBITE(args.testdir, input_transform, target_transform))
train_loader = DataLoader(
MNIBITE(args.traindir, input_transform, target_transform))
for epoch in range(args.num_epochs):
train_epoch(args, epoch, model, train_loader, crit1, crit2, optimizer,
trainer)
trainer.vis_losses(epoch)
trainer.reset()
test_epoch(args, epoch, model, test_loader, crit1, crit2, trainer)
trainer.vis_losses(epoch)
trainer.reset()
if args.save_epochs > 0 and epoch % args.save_epochs == 0:
stamp = time.strftime("%Y%m%d-%H%M%S")
torch.save(model.state_dict(), f'{args.name}-{stamp}.pth')
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
# defining net, loss and optimizer
l2_loss = MSELoss()
net = DNCNN(num_channels=64)
# check if saved file of model exists, if not continue without loading
if os.path.isfile(net_path) and load_net:
best_loss_score = load_checkpoint(net, net_path)
# setting to cuda
if cuda >= 0:
net.cuda(cuda)
l2_loss.cuda(cuda)
optimizer = optim.Adam(net.parameters(), lr=learning_rate)
if train_net:
# variables to keep the progress
loss_history = []
counter = []
# training
net.train()
l2_loss.train()
for epoch in range(num_epochs):
tic = time.clock()
# train routine - maybe i want to define a function for it
num_input = Sx_tr.shape[-1]
num_output = y_tr.shape[-1]
num_neuron = 30
# define networks
model = Sequential(\
Linear(num_input, num_neuron), LeakyReLU(),\
Linear(num_neuron, num_neuron), LeakyReLU(),\
Linear(num_neuron, num_output))
optimizer = Adam(model.parameters())
criterion = MSELoss()
#----------------------------------------------------------------------------------------------
# use cuda
model = model.cuda()
criterion = criterion.cuda()
#----------------------------------------------------------------------------------------------
# Number of signals to use in each gradient descent step (batch).
batch_size = training_size // 10
# Number of epochs.
num_epochs = 1e4
# Learning rate for Adam.
lr = 1e-4
#----------------------------------------------------------------------------------------------
# break into batches
nsamples = Sx_tr.shape[0]
nbatches = nsamples // batch_size
class Trainer:
def __init__(self,
num_epochs: int = 500,
cuda: bool = True,
continue_from: int = 0):
self.features = 16
self.instances = 64
self.classes = 2
self.z_size = 100
self.batch_size = 100
self.workers = 5
self.num_epochs = num_epochs
self.cuda = cuda
self.log_step = 10
self.log_step_print = 50
self.save_period = 5
self.continue_from = continue_from
self.models_path = "./models_grid"
self.lambdas = LambdaFeaturesCollector(self.features, self.instances)
self.metas = MetaFeaturesCollector(self.features, self.instances)
self.data_loader = get_loader(
f"../processed_data/processed_{self.features}_{self.instances}_{self.classes}/",
self.features, self.instances, self.classes, self.metas,
self.lambdas, self.batch_size, self.workers)
self.test_loader = get_loader(f"../processed_data/test/",
16,
64,
2,
self.metas,
self.lambdas,
228,
5,
train_meta=False)
if continue_from == 0:
self.generator = Generator(self.features,
self.instances, self.classes,
self.metas.getLength(), self.z_size)
self.discriminator = Discriminator(self.features, self.instances,
self.classes,
self.metas.getLength(),
self.lambdas.getLength())
else:
self.generator = Generator(self.features,
self.instances, self.classes,
self.metas.getLength(), self.z_size)
self.generator.load_state_dict(
torch.load(
f'{self.models_path}/generator-{self.features}_{self.instances}_{self.classes}-{continue_from}.pkl'
))
self.generator.eval()
self.discriminator = Discriminator(self.features, self.instances,
self.classes,
self.metas.getLength(),
self.lambdas.getLength())
self.discriminator.load_state_dict(
torch.load(
f'{self.models_path}/discriminator-{self.features}_{self.instances}_{self.classes}-{continue_from}.pkl'
))
self.discriminator.eval()
if self.cuda:
self.generator.cuda()
if self.cuda:
self.discriminator.cuda()
self.lr = 0.0002
self.beta1 = 0.5
self.beta2 = 0.999
self.g_optimizer = optim.Adam(self.generator.parameters(), self.lr,
[self.beta1, self.beta2])
self.d_optimizer = optim.Adam(self.discriminator.parameters(), self.lr,
[self.beta1, self.beta2])
self.cross_entropy = BCEWithLogitsLoss()
if self.cuda:
self.cross_entropy.cuda()
self.mse = MSELoss()
if self.cuda:
self.mse.cuda()
def to_variable(self, x):
if self.cuda:
x = x.cuda()
return Variable(x)
def getDistance(self, x: torch.Tensor, y: torch.Tensor) -> [float]:
x_in = np.squeeze(x.cpu().detach().numpy())
y_in = np.squeeze(y.cpu().detach().numpy())
results = []
for (xx, yy) in zip(x_in, y_in):
try:
V = np.cov(np.array([xx, yy]).T)
V[np.diag_indices_from(V)] += 0.1
IV = np.linalg.inv(V)
D = mahalanobis(xx, yy, IV)
except:
D = 0.0
results.append(D)
return results
def getMeta(self, data_in: torch.Tensor):
meta_list = []
for data in data_in:
meta_list.append(self.metas.getShort(data.cpu().detach().numpy()))
result = torch.stack(meta_list)
return self.to_variable(
result.view((result.size(0), result.size(1), 1, 1)))
def getLambda(self, data_in: torch.Tensor):
lamba_list = []
for data in data_in:
lamba_list.append(self.lambdas.get(data.cpu().detach().numpy()))
result = torch.stack(lamba_list)
return self.to_variable(result)
def train(self):
total_steps = len(self.data_loader)
logging.info(f'Starting training...')
for epoch in range(self.continue_from, self.num_epochs):
loss = []
q = 0
for i, data in enumerate(self.test_loader):
q += 1
dataset = self.to_variable(data[0])
metas = self.to_variable(data[1])
lambdas = self.to_variable(data[2])
real_outputs = self.discriminator(dataset, metas)
d_real_labels_loss = self.mse(real_outputs[:, 1:], lambdas)
loss.append(d_real_labels_loss.cpu().detach().numpy())
logging.info(f'{epoch}d:{np.mean(loss)}')
results = []
q = 0
for i, data in enumerate(self.test_loader):
q += 1
metas = self.to_variable(data[1])
batch_size = data[0].size(0)
noise = torch.randn(batch_size, 100)
noise = noise.view((noise.size(0), noise.size(1), 1, 1))
noise = self.to_variable(noise)
fake_data = self.generator(noise, metas)
fake_metas = self.getMeta(fake_data)
results.extend(self.mse(fake_metas, metas))
logging.info(f'{epoch}g:{np.mean(np.array(results))}')
q = 0
for i, data in enumerate(self.data_loader):
q += 1
dataset = self.to_variable(data[0])
metas = self.to_variable(data[1])
lambdas = self.to_variable(data[2])
batch_size = data[0].size(0)
noise = torch.randn(batch_size, self.z_size)
noise = noise.view((noise.size(0), noise.size(1), 1, 1))
noise = self.to_variable(noise)
zeros = torch.zeros([batch_size, 1], dtype=torch.float32)
zeros = self.to_variable(zeros)
ones = torch.ones([batch_size, 1], dtype=torch.float32)
ones = self.to_variable(ones)
# Get D on real
real_outputs = self.discriminator(dataset, metas)
d_real_labels_loss = self.mse(real_outputs[:, 1:], lambdas)
d_real_rf_loss = self.mse(real_outputs[:, :1], zeros)
d_real_loss = d_real_labels_loss + 0.7 * d_real_rf_loss
# Get D on fake
fake_data = self.generator(noise, metas)
fake_data_metas = self.getMeta(fake_data)
fake_outputs = self.discriminator(fake_data, fake_data_metas)
fake_lambdas = self.getLambda(fake_data)
d_fake_labels_loss = self.mse(fake_outputs[:, 1:],
fake_lambdas)
d_fake_rf_loss = self.mse(fake_outputs[:, :1], ones)
d_fake_loss = 0.7 * d_fake_rf_loss + 0.6 * d_fake_labels_loss
# Train D
d_loss = d_real_loss + 0.8 * d_fake_loss
self.generator.zero_grad()
self.discriminator.zero_grad()
d_loss.backward()
self.d_optimizer.step()
# Get D on fake
noise = torch.randn(batch_size, self.z_size)
noise = noise.view(noise.size(0), noise.size(1), 1, 1)
noise = self.to_variable(noise)
fake_data = self.generator(noise, metas)
fake_outputs = self.discriminator(fake_data, metas)
g_fake_rf_loss = self.mse(fake_outputs[:, :1], zeros)
fake_metas = self.getMeta(fake_data)
g_fake_meta_loss = self.mse(fake_metas, metas)
g_loss = 0.7 * g_fake_rf_loss + g_fake_meta_loss
# Train G
self.generator.zero_grad()
self.discriminator.zero_grad()
g_loss.backward()
self.g_optimizer.step()
# logging
#if (q + 1) % self.log_step == 0:
log = (
f'[[{epoch},{i}],[{d_real_rf_loss},{d_real_labels_loss},{d_fake_rf_loss},{d_fake_labels_loss}],[{g_fake_rf_loss},{g_fake_meta_loss}]]'
)
logging.info(log)
#if (q + 1) % self.log_step_print == 0:
print((
f'[{datetime.now()}] Epoch[{epoch}/{self.num_epochs}], Step[{q}/{total_steps}],'
f' D_losses: [{d_real_rf_loss}|{d_real_labels_loss}|{d_fake_rf_loss}|{d_fake_labels_loss}], '
f'G_losses:[{g_fake_rf_loss}|{g_fake_meta_loss}]'))
# saving
if (epoch + 1) % self.save_period == 0:
done_data_str_path = Path(self.models_path)
done_data_str_path.mkdir(parents=True, exist_ok=True)
g_path = os.path.join(
self.models_path,
f'generator-{self.features}_{self.instances}_{self.classes}-{epoch + 1}.pkl'
)
d_path = os.path.join(
self.models_path,
f'discriminator-{self.features}_{self.instances}_{self.classes}-{epoch + 1}.pkl'
)
torch.save(self.generator.state_dict(), g_path)
torch.save(self.discriminator.state_dict(), d_path)
class YOLOLayer(Module):
"""Detection layer"""
def __init__(self, anchors, num_classes, img_dim):
super(YOLOLayer, self).__init__()
self.anchors = anchors
self.num_anchors = len(anchors)
self.num_classes = num_classes
self.bbox_attrs = 5 + num_classes
self.img_dim = img_dim
self.ignore_thres = 0.5
self.lambda_coord = 1
self.mse_loss = MSELoss()
self.bce_loss = BCELoss()
def forward(self, x, targets=None):
bs = x.size(0)
g_dim = x.size(2)
stride = self.img_dim / g_dim
# Tensors for cuda support
FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
prediction = x.view(bs, self.num_anchors, self.bbox_attrs, g_dim, g_dim).permute(0, 1, 3, 4, 2).contiguous()
# Get outputs
x = torch.sigmoid(prediction[..., 0]) # Center x
y = torch.sigmoid(prediction[..., 1]) # Center y
w = prediction[..., 2] # Width
h = prediction[..., 3] # Height
conf = torch.sigmoid(prediction[..., 4]) # Conf
pred_cls = torch.sigmoid(prediction[..., 5:]) # Cls pred.
# Calculate offsets for each grid
grid_x = torch.linspace(0, g_dim-1, g_dim).repeat(g_dim,1).repeat(bs*self.num_anchors, 1, 1).view(x.shape).type(FloatTensor)
grid_y = torch.linspace(0, g_dim-1, g_dim).repeat(g_dim,1).t().repeat(bs*self.num_anchors, 1, 1).view(y.shape).type(FloatTensor)
scaled_anchors = [(a_w / stride, a_h / stride) for a_w, a_h in self.anchors]
anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, g_dim*g_dim).view(w.shape)
anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, g_dim*g_dim).view(h.shape)
# Add offset and scale with anchors
pred_boxes = FloatTensor(prediction[..., :4].shape)
pred_boxes[..., 0] = x.data + grid_x
pred_boxes[..., 1] = y.data + grid_y
pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
pred_boxes[..., 3] = torch.exp(h.data) * anchor_h
# Training
if targets is not None:
if x.is_cuda:
self.mse_loss = self.mse_loss.cuda()
self.bce_loss = self.bce_loss.cuda()
nGT, nCorrect, mask, conf_mask, tx, ty, tw, th, tconf, tcls = build_targets(pred_boxes.cpu().data,
targets.cpu().data,
scaled_anchors,
self.num_anchors,
self.num_classes,
g_dim,
self.ignore_thres,
self.img_dim)
nProposals = int((conf > 0.25).sum().item())
recall = float(nCorrect / nGT) if nGT else 1
# Handle masks
mask = Variable(mask.type(FloatTensor))
cls_mask = Variable(mask.unsqueeze(-1).repeat(1, 1, 1, 1, self.num_classes).type(FloatTensor))
conf_mask = Variable(conf_mask.type(FloatTensor))
# Handle target variables
tx = Variable(tx.type(FloatTensor), requires_grad=False)
ty = Variable(ty.type(FloatTensor), requires_grad=False)
tw = Variable(tw.type(FloatTensor), requires_grad=False)
th = Variable(th.type(FloatTensor), requires_grad=False)
tconf = Variable(tconf.type(FloatTensor), requires_grad=False)
tcls = Variable(tcls.type(FloatTensor), requires_grad=False)
# Mask outputs to ignore non-existing objects
loss_x = self.lambda_coord * self.bce_loss(x * mask, tx * mask)
loss_y = self.lambda_coord * self.bce_loss(y * mask, ty * mask)
loss_w = self.lambda_coord * self.mse_loss(w * mask, tw * mask) / 2
loss_h = self.lambda_coord * self.mse_loss(h * mask, th * mask) / 2
loss_conf = self.bce_loss(conf * conf_mask, tconf * conf_mask)
loss_cls = self.bce_loss(pred_cls * cls_mask, tcls * cls_mask)
loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls
return loss, loss_x.item(), loss_y.item(), loss_w.item(), loss_h.item(), loss_conf.item(), loss_cls.item(), recall
else:
# If not in training phase return predictions
output = torch.cat((pred_boxes.view(bs, -1, 4) * stride, conf.view(bs, -1, 1), pred_cls.view(bs, -1, self.num_classes)), -1)
return output.data
def main(opt):
cuda = True if torch.cuda.is_available() else False
device = torch.device('cuda') if cuda else torch.device('cpu')
# Define dataset.
noise_args = get_noise_args(opt.noise, L=opt.L)
transform = transforms.Compose(
[transforms.RandomCrop(opt.crop_size),
transforms.ToTensor(),
MultiplicativeNoise(**noise_args) # returns (noisy, clean) tuple
]
)
# dataset returns (noisy, clean) tuple
dataset = CMCropsDataset(opt.data_root, only_F=True, transform_F=transform)
train_size = int(0.9 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
train_dataloader = DataLoader(train_dataset, batch_size=opt.batch_size, num_workers=4)
val_dataloader = DataLoader(val_dataset, batch_size=opt.batch_size, num_workers=4)
# Define model and loss criterion.
model = get_model(opt.model, opt.layers, opt.filters, opt.filters_size,
apply_sigmoid=not opt.no_sigmoid)
if opt.criterion == 'mse':
criterion = MSELoss()
elif opt.criterion == 'l1':
criterion = L1Loss()
if cuda:
model = model.cuda()
criterion = criterion.cuda()
# Define optimizer.
if opt.optim == 'adam':
optimizer = Adam(params=model.parameters(), lr=opt.lr)
else:
raise NotImplementedError(opt.optim + ' optimizer is not supported.')
# Training process.
loss_hist = list() # list of (epoch, train loss)
loss_hist_eval = list() # list of (epoch, validation loss)
ssim_hist_eval = list() # list of (epoch, validation SSIM)
for epoch in range(opt.epochs):
# TRAINING.
model.train()
input_and_target = enumerate(train_dataloader)
if opt.verbose:
print('Epoch {} of {}'.format(epoch, opt.epochs - 1))
input_and_target = tqdm.tqdm(input_and_target, total=len(train_dataloader))
med_loss = 0
for i, (x_batch, target_batch) in input_and_target:
x_batch, target_batch = x_batch.float().to(device), target_batch.float().to(device)
optimizer.zero_grad()
output_batch = model(x_batch)
loss = criterion(output_batch, target_batch)
loss.backward()
optimizer.step()
med_loss += loss.data.cpu().numpy()
if opt.verbose:
input_and_target.set_description('Train loss = {0:.3f}'.format(loss))
loss_hist.append((epoch, med_loss / (i + 1)))
# VALIDATION.
if opt.verbose:
print('Validation:')
model.eval()
with torch.no_grad():
input_and_target = enumerate(val_dataloader)
if opt.verbose:
input_and_target = tqdm.tqdm(input_and_target, total=len(val_dataloader))
med_loss_eval = 0
med_ssim_eval = 0
for i, (x_batch, target_batch) in input_and_target:
x_batch, target_batch = x_batch.float().to(device), target_batch.float().to(device)
output_batch = model(x_batch)
loss = criterion(output_batch, target_batch)
med_loss_eval += loss.data.cpu().numpy()
prev_loss_eval = criterion(x_batch, target_batch).data.cpu().numpy()
ssim_input = compute_ssim(x_batch, target_batch)
ssim_output = compute_ssim(output_batch, target_batch, opt.median)
med_ssim_eval += ssim_output
if opt.verbose:
input_and_target.set_description(
'Validation: '
+ 'Output loss = {0:.3f}'.format(loss)
+ ' Input loss = {0:.3f}'.format(prev_loss_eval)
+ ' Input SSIM = {0:.3f}'.format(ssim_input)
+ ' Output SSIM = {0:.3f}'.format(ssim_output))
loss_hist_eval.append((epoch, med_loss_eval / (i + 1)))
ssim_hist_eval.append((epoch, med_ssim_eval / (i + 1)))
write_lc_step(epoch, loss_hist[-1][1], loss_hist_eval[-1][1], ssim_hist_eval[-1][1])
torch.save(model.state_dict(), os.path.join(opt.output, '{}_latest.h5'.format(opt.model)))
if epoch % opt.save_period == 0:
torch.save(model.state_dict(), os.path.join(opt.output, '{}_epoch{}.h5'.format(opt.model, epoch)))
torch.save(model.state_dict(), os.path.join(opt.output, 'model_latest.h5'))
with open(os.path.join(opt.output, 'train_loss_hist.pkl'), 'wb') as f:
pickle.dump(loss_hist, f, pickle.HIGHEST_PROTOCOL)
with open(os.path.join(opt.output, 'valid_loss_hist.pkl'), 'wb') as f:
pickle.dump(loss_hist_eval, f, pickle.HIGHEST_PROTOCOL)
def __init__(self,
epochs=50,
h1=128,
lr=0.07,
data_file="data/Compiled_Data.csv",
save=True):
self.data_file = data_file
self.h1 = h1
self.filename = f'Outputs/{__class__.__name__}/{epochs}_{h1}_{lr}'
self.lr = lr
self.MODEL = helpdesk.Model4
# saving the model .pth and the image .png
self.SAVE = save
# defining the number of epochs
self.epochs = epochs
# empty list to store training losses
self.train_losses = []
# empty list to store validation losses
self.val_losses = []
# plotION default off
self.plotion = False
# loading dataset
data = pd.read_csv(data_file).to_numpy()
random.shuffle(data)
train_x = data[:, 1:]
train_y = data[:, 0].reshape((-1, 1))
# create validation set
train_x, val_x, train_y, val_y = train_test_split(train_x,
train_y,
test_size=0.1)
print((train_x.shape, train_y.shape), (val_x.shape, val_y.shape))
# converting into torch format
self.train_x = torch.from_numpy(train_x)
self.train_y = torch.from_numpy(train_y)
self.val_x = torch.from_numpy(val_x)
self.val_y = torch.from_numpy(val_y)
self._in = self.train_x.shape[1]
self._out = self.train_y.shape[1]
# defining the model
print('Build Model')
model = helpdesk.Model4(_in=self._in, H=h1, _out=self._out)
model.double()
# defining the optimizer
optimizer = Adam(model.parameters(), lr=lr)
# defining the loss function
criterion = MSELoss()
# checking if GPU is available
if torch.cuda.is_available():
model = model.cuda()
criterion = criterion.cuda()
self.model = model
self.optimizer = optimizer
self.criterion = criterion
self.info = str(self)
print(model)
class TrainerCNN:
def __init__(self,
num_epochs: int = 20,
cuda: bool = True,
continue_from: int = 0):
self.features = 16
self.instances = 64
self.classes = 2
self.z_size = 100
self.batch_size = 100
self.workers = 4
self.num_epochs = num_epochs
self.cuda = cuda
self.log_step = 10
self.log_step_print = 50
self.save_period = 5
self.continue_from = continue_from
self.models_path = "./cnn1206"
self.lambdas = LambdaFeaturesCollector(self.features, self.instances)
self.metas = MetaZerosCollector(self.features, self.instances)
self.data_loader = get_loader(
f"../processed_data/processed_{self.features}_{self.instances}_{self.classes}/",
self.features, self.instances, self.classes, self.metas,
self.lambdas, self.batch_size, self.workers)
self.test_loader = get_loader(f"../processed_data/test/",
16,
64,
2,
self.metas,
self.lambdas,
147,
self.workers,
train_meta=False)
if continue_from == 0:
self.discriminator = Discriminator(self.features, self.instances,
self.classes,
self.metas.getLength(),
self.lambdas.getLength())
else:
self.discriminator = Discriminator(self.features, self.instances,
self.classes,
self.metas.getLength(),
self.lambdas.getLength())
self.discriminator.load_state_dict(
torch.load(
f'{self.models_path}/discriminator-{self.features}_{self.instances}_{self.classes}-{continue_from}.pkl'
))
self.discriminator.eval()
if self.cuda:
self.discriminator.cuda()
self.lr = 0.0002
self.beta1 = 0.5
self.beta2 = 0.999
self.d_optimizer = optim.Adam(self.discriminator.parameters(), self.lr,
[self.beta1, self.beta2])
self.cross_entropy = BCEWithLogitsLoss()
if self.cuda:
self.cross_entropy.cuda()
self.mse = MSELoss()
if self.cuda:
self.mse.cuda()
def to_variable(self, x):
if self.cuda:
x = x.cuda()
return Variable(x)
def train(self):
total_steps = len(self.data_loader)
logging.info(f'Starting training...')
for epoch in range(self.continue_from, self.num_epochs):
loss = []
for i, data in enumerate(self.test_loader):
dataset = self.to_variable(data[0])
metas = self.to_variable(data[1])
lambdas = self.to_variable(data[2])
real_outputs = self.discriminator(dataset, metas)
d_real_labels_loss = self.mse(real_outputs[:, 1:], lambdas)
loss.append(d_real_labels_loss.cpu().detach().numpy())
logging.info(f'{epoch}d:{np.mean(loss)}')
for i, data in enumerate(self.data_loader):
dataset = self.to_variable(data[0])
metas = self.to_variable(data[1])
lambdas = self.to_variable(data[2])
# Get D on real
real_outputs = self.discriminator(dataset, metas)
d_real_labels_loss = self.mse(real_outputs[:, 1:], lambdas)
d_real_loss = d_real_labels_loss
# Train D
d_loss = d_real_loss
self.discriminator.zero_grad()
d_loss.backward()
self.d_optimizer.step()
# # logging
# if (i + 1) % self.log_step == 0:
# log = (
# f'[[{epoch},{i}],[{d_real_labels_loss}]'
# )
# logging.info(log)
if (i + 1) % self.log_step_print == 0:
print((
f'[{datetime.now()}] Epoch[{epoch}/{self.num_epochs}], Step[{i}/{total_steps}],'
f' D_losses: [{d_real_labels_loss}], '))
# saving
if (epoch + 1) % self.save_period == 0:
done_data_str_path = Path(self.models_path)
done_data_str_path.mkdir(parents=True, exist_ok=True)
d_path = os.path.join(
self.models_path,
f'discriminator-{self.features}_{self.instances}_{self.classes}-{epoch + 1}.pkl'
)
torch.save(self.discriminator.state_dict(), d_path)
def train_mnist():
ag = []
nb_class = 10
img_size = 28
n = 64
f = 7
n_m = 12
d = 2
nb_action = 4
batch_size = 64
t = 7
nr = 1
cuda = True
#m = ModelsUnion(n, f, n_m, d, nb_action, nb_class, test_mnist())
m = ModelsUnion(n, f, n_m, d, nb_action, nb_class)
a1 = Agent(ag, m, n, f, n_m, img_size, nb_action, batch_size, obs_MNIST,
trans_MNIST)
a2 = Agent(ag, m, n, f, n_m, img_size, nb_action, batch_size, obs_MNIST,
trans_MNIST)
a3 = Agent(ag, m, n, f, n_m, img_size, nb_action, batch_size, obs_MNIST,
trans_MNIST)
ag.append(a1)
ag.append(a2)
ag.append(a3)
if cuda:
for a in ag:
a.cuda()
sm = Softmax(dim=-1)
criterion = MSELoss()
if cuda:
criterion.cuda()
params = []
for net in m.get_networks():
if cuda:
net.cuda()
params += list(net.parameters())
optim = th.optim.Adam(params, lr=1e-3)
nb_epoch = 10
(x_train, y_train), (x_valid, y_valid), (x_test, y_test) = load_mnist()
x_train, y_train = x_train[:10000], y_train[:10000]
nb_batch = ceil(x_train.size(0) / batch_size)
loss_v = []
acc = []
for e in range(nb_epoch):
sum_loss = 0
for net in m.get_networks():
net.train()
grad_norm_cnn = []
grad_norm_pred = []
random_walk = e < 5
for i in tqdm(range(nb_batch)):
i_min = i * batch_size
i_max = (i + 1) * batch_size
i_max = i_max if i_max < x_train.size(0) else x_train.size(0)
losses = []
for k in range(nr):
x, y = x_train[i_min:i_max, :, :], y_train[i_min:i_max]
if cuda:
x, y = x.cuda(), y.cuda()
pred, log_probas = step(ag, x, t, sm, cuda, random_walk,
nb_class)
# Sum on agent dimension
proba_per_image = log_probas.sum(dim=0)
y_eye = th.eye(nb_class)[y]
if cuda:
y_eye = y_eye.cuda()
r = -criterion(pred, y_eye)
# Mean on image batch
l = (proba_per_image * r.detach() + r).mean(dim=0).view(-1)
losses.append(l)
loss = -th.cat(losses).sum() / nr
optim.zero_grad()
loss.backward()
optim.step()
sum_loss += loss.item()
grad_norm_cnn.append(
m.get_networks()[0].seq_lin[0].weight.grad.norm())
grad_norm_pred.append(
m.get_networks()[-1].seq_lin[0].weight.grad.norm())
sum_loss /= nb_batch
print("Epoch %d, loss = %f" % (e, sum_loss))
print("grad_cnn_norm_mean = %f, grad_pred_norm_mean = %f" %
(sum(grad_norm_cnn) / len(grad_norm_cnn),
sum(grad_norm_pred) / len(grad_norm_pred)))
print("CNN_el = %d, Pred_el = %d" %
(m.get_networks()[0].seq_lin[0].weight.grad.nelement(),
m.get_networks()[-1].seq_lin[0].weight.grad.nelement()))
nb_correct = 0
nb_batch_valid = ceil(x_valid.size(0) / batch_size)
for net in m.get_networks():
net.eval()
with th.no_grad():
for i in tqdm(range(nb_batch_valid)):
i_min = i * batch_size
i_max = (i + 1) * batch_size
i_max = i_max if i_max < x_valid.size(0) else x_valid.size(0)
x, y = x_valid[
i_min:i_max, :, :].cuda(), y_valid[i_min:i_max].cuda()
pred, proba = step(ag, x, t, sm, cuda, random_walk, nb_class)
nb_correct += (pred.argmax(dim=1) == y).sum().item()
nb_correct /= x_valid.size(0)
acc.append(nb_correct)
loss_v.append(sum_loss)
print("Epoch %d, accuracy = %f" % (e, nb_correct))
plt.plot(acc, "b", label="accuracy")
plt.plot(loss_v, "r", label="criterion value")
plt.xlabel("Epoch")
plt.title("MARL Classification f=%d, n=%d, n_m=%d, d=%d, T=%d" %
(f, n, n_m, d, t))
plt.legend()
plt.show()
viz(ag, x_test[randint(0, x_test.size(0) - 1)], t, sm, f)
shuffle=True)
test_data_loader = DataLoader(dataset=test_set,
num_workers=2,
batch_size=BATCH_SIZE,
shuffle=True)
test_input, test_target = test_data_loader.__iter__().__next__()
real_a = torch.FloatTensor(BATCH_SIZE, IMAGE_CHANNEL, IMAGE_SIZE, IMAGE_SIZE)
real_b = torch.FloatTensor(BATCH_SIZE, OUTPUT_CHANNEL, IMAGE_SIZE, IMAGE_SIZE)
if GPU_NUMS > 1:
Net_G = Net_G.cuda()
Net_D = Net_D.cuda()
lossGAN = lossGAN.cuda()
lossL1 = lossL1.cuda()
lossMSE = lossMSE.cuda()
real_a = Variable(real_a.cuda() if GPU_NUMS > 1 else real_a)
real_b = Variable(real_b.cuda() if GPU_NUMS > 1 else real_b)
bar = ProgressBar(EPOCHS, len(train_data_loader), "D loss:%.3f;G loss:%.3f")
for epoch in range(EPOCHS):
for iteration, batch in enumerate(train_data_loader, 1):
real_a_cpu, real_b_cpu = batch[0], batch[1]
real_a.data.resize_(real_a_cpu.size()).copy_(real_a_cpu)
real_b.data.resize_(real_b_cpu.size()).copy_(real_b_cpu)
fake_b = Net_G(real_a)
optimizer_D.zero_grad()
# train with fake
