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