def __init__(self, hyperparameters, gpu_id): super(Discriminator, self).__init__() self.hyp = hyperparameters self.gpu_id = gpu_id self.w_loss_meter = MovingAverageValueMeter(5) self.d_loss_meter = MovingAverageValueMeter(5) self.r_loss_meter = MovingAverageValueMeter(5) self.f_loss_meter = MovingAverageValueMeter(5) self.gp_loss_meter = MovingAverageValueMeter(5) # Architecture self.lab0 = nn.ConvTranspose2d(1, self.hyp['q1'], (1, 55), bias=False) self.conv0 = nn.Conv2d(1, self.hyp['q2'], (55, 1), bias=False) self.nonlin0 = nn.Sequential(*[ nn.BatchNorm2d(self.hyp['q1'] + self.hyp['q2']), nn.LeakyReLU(self.hyp['lrelu_d']) ] if self.hyp['bd0'] else [ nn.LeakyReLU(self.hyp['lrelu_d']), ]) self.conv1 = nn.Conv2d(self.hyp['q1'] + self.hyp['q2'], self.hyp['q3'], (1, 55), bias=False) self.nonlin1 = nn.Sequential(*[ nn.BatchNorm2d(self.hyp['q3']), nn.LeakyReLU(self.hyp['lrelu_d']) ] if self.hyp['bd1'] else [ nn.LeakyReLU(self.hyp['lrelu_d']), ]) self.fc = nn.Linear(self.hyp['q3'], 1, bias=False) self.cuda(self.gpu_id) opt_param_list = [{ 'params': [ param for name, param in self.named_parameters() if 'lab0' not in name ] }, { 'params': self.lab0.parameters(), 'lr': 1 * self.hyp['lr_d'] }] self.optimizer = torch.optim.Adam(opt_param_list, lr=self.hyp['lr_d'], betas=(self.hyp['b1_d'], self.hyp['b2_d']), weight_decay=self.hyp['wd_d']) # rand init for m in self.modules(): if isinstance(m, nn.ConvTranspose2d): torch.nn.init.kaiming_normal_(m.weight, a=self.hyp['lrelu_d'], nonlinearity='leaky_relu') if not m.bias is None: torch.nn.init.constant_(m.bias, 0) if isinstance(m, nn.Conv2d): torch.nn.init.kaiming_normal_(m.weight, a=self.hyp['lrelu_d'], nonlinearity='leaky_relu') if not m.bias is None: torch.nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): torch.nn.init.constant_(m.weight, 1) torch.nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): torch.nn.init.kaiming_normal_(m.weight) if not m.bias is None: torch.nn.init.constant_(m.bias, 0)
def main(config, cuda): cuda = cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") if cuda: current_device = torch.cuda.current_device() print("Running on", torch.cuda.get_device_name(current_device)) else: print("Running on CPU") # Configuration CONFIG = Dict(yaml.load(open(config))) dataset = get_dataset(CONFIG.DATASET)( data_path=CONFIG.ROOT, crop_size=256, scale=(0.6, 0.8, 1., 1.2, 1.4), rotation=15, flip=True, mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R), ) """ # Dataset 10k or 164k dataset = get_dataset(CONFIG.DATASET)( root=CONFIG.ROOT, split=CONFIG.SPLIT.TRAIN, base_size=513, crop_size=CONFIG.IMAGE.SIZE.TRAIN, mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R), warp=CONFIG.WARP_IMAGE, scale=(0.5, 0.75, 1.0, 1.25, 1.5), flip=True, ) """ # DataLoader loader = torch.utils.data.DataLoader( dataset=dataset, batch_size=CONFIG.BATCH_SIZE.TRAIN, num_workers=CONFIG.NUM_WORKERS, shuffle=True, ) loader_iter = iter(loader) # Model model = DeepLabV3Plus_ResNet101_MSC(n_classes=CONFIG.N_CLASSES) state_dict = torch.load(CONFIG.INIT_MODEL) model.load_state_dict(state_dict, strict=False) # Skip "aspp" layer model = nn.DataParallel(model) model.to(device) for name, param in model.named_parameters(): if param.requires_grad: print(name) # Optimizer optimizer = torch.optim.Adam( params=get_params(model.module), lr=CONFIG.LR, weight_decay=CONFIG.WEIGHT_DECAY, ) """ # Optimizer optimizer = torch.optim.SGD( # cf lr_mult and decay_mult in train.prototxt params=[ { "params": get_params(model.module, key="1x"), "lr": CONFIG.LR, "weight_decay": CONFIG.WEIGHT_DECAY, }, { "params": get_params(model.module, key="10x"), "lr": 10 * CONFIG.LR, "weight_decay": CONFIG.WEIGHT_DECAY, }, { "params": get_params(model.module, key="20x"), "lr": 20 * CONFIG.LR, "weight_decay": 0.0, }, ], momentum=CONFIG.MOMENTUM, ) """ # Loss definition criterion = CrossEntropyLoss2d(ignore_index=CONFIG.IGNORE_LABEL) criterion.to(device) max_pooling_loss = MaxPoolingLoss(ratio=0.3, p=1.7, reduce=True) # TensorBoard Logger writer = SummaryWriter(CONFIG.LOG_DIR) loss_meter = MovingAverageValueMeter(20) model.train() model.module.scale.freeze_bn() for iteration in tqdm( range(1, CONFIG.ITER_MAX + 1), total=CONFIG.ITER_MAX, leave=False, dynamic_ncols=True, ): """ # Set a learning rate poly_lr_scheduler( optimizer=optimizer, init_lr=CONFIG.LR, iter=iteration - 1, lr_decay_iter=CONFIG.LR_DECAY, max_iter=CONFIG.ITER_MAX, power=CONFIG.POLY_POWER, ) """ # Clear gradients (ready to accumulate) optimizer.zero_grad() iter_loss = 0 for i in range(1, CONFIG.ITER_SIZE + 1): try: images, labels = next(loader_iter) except: loader_iter = iter(loader) images, labels = next(loader_iter) images = images.to(device) labels = labels.to(device).unsqueeze(1).float() # Propagate forward logits = model(images) # Loss loss = 0 for logit in logits: # Resize labels for {100%, 75%, 50%, Max} logits labels_ = F.interpolate(labels, logit.shape[2:], mode="nearest") labels_ = labels_.squeeze(1).long() # Compute NLL and MPL nll_loss = criterion(logit, labels_) # loss += nll_loss loss += max_pooling_loss(nll_loss) # Backpropagate (just compute gradients wrt the loss) loss /= float(CONFIG.ITER_SIZE) loss.backward() iter_loss += float(loss) loss_meter.add(iter_loss) # Update weights with accumulated gradients optimizer.step() if iteration % CONFIG.ITER_TB == 0: writer.add_scalar("train_loss", loss_meter.value()[0], iteration) for i, o in enumerate(optimizer.param_groups): writer.add_scalar("train_lr_group{}".format(i), o["lr"], iteration) gt_viz, images_viz, predicts_viz = make_vizs( images, labels_, logits, (CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R)) writer.add_image("gt/images", torch.from_numpy(images_viz[0]), iteration) writer.add_image("gt/labels", torch.from_numpy(gt_viz[0]), iteration) for i, predict_viz in enumerate(predicts_viz): writer.add_image("predict/" + str(i), torch.from_numpy(predict_viz[0]), iteration) if False: # This produces a large log file for name, param in model.named_parameters(): name = name.replace(".", "/") writer.add_histogram(name, param, iteration, bins="auto") if param.requires_grad: writer.add_histogram(name + "/grad", param.grad, iteration, bins="auto") # Save a model if iteration % CONFIG.ITER_SAVE == 0: torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, "checkpoint_{}.pth".format(iteration)), ) # Save a model (short term) if iteration % 100 == 0: torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, "checkpoint_current.pth"), ) torch.save(model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, "checkpoint_final.pth"))
def __init__(self, hyperparameters, log_dir, gpu_id): super(Generator, self).__init__() self.hyp = hyperparameters print(hyperparameters) self.gpu_id = gpu_id self.noise_dim = self.hyp['noise_dim'] self.vis_noise = torch.randn(1, self.hyp['noise_dim']).cuda( self.gpu_id).requires_grad_(False) self.g_loss_meter = MovingAverageValueMeter(5) self.log_dir = log_dir # Architecture: self.lab0 = nn.Linear(1, self.hyp['p1'], bias=False) self.fc0 = nn.Linear(self.noise_dim, self.hyp['p2'], bias=False) self.nonlin0 = nn.Sequential(*[ nn.BatchNorm2d(self.hyp['p1'] + self.hyp['p2']), nn.LeakyReLU(self.hyp['lrelu_g']) ] if self.hyp['bg0'] else [ nn.LeakyReLU(self.hyp['lrelu_g']), ]) self.conv1 = nn.ConvTranspose2d(self.hyp['p1'] + self.hyp['p2'], self.hyp['p3'], (1, 55), bias=True) self.nonlin1 = nn.Sequential(*[ nn.BatchNorm2d(self.hyp['p3']), nn.LeakyReLU(self.hyp['lrelu_g']) ] if self.hyp['bg1'] else [ nn.LeakyReLU(self.hyp['lrelu_g']), ]) self.conv2 = nn.ConvTranspose2d(self.hyp['p3'], 1, (55, 1), bias=True) self.sigmoid = nn.Tanh() self.cuda(self.gpu_id) opt_param_list = [{ 'params': [ param for name, param in self.named_parameters() if 'lab0' not in name ] }, { 'params': self.lab0.parameters(), 'lr': 1 * self.hyp['lr_g'] }] self.optimizer = torch.optim.Adam(opt_param_list, lr=self.hyp['lr_g'], betas=(self.hyp['b1_g'], self.hyp['b2_g']), weight_decay=self.hyp['wd_g']) # rand init for m in self.modules(): if isinstance(m, nn.ConvTranspose2d): torch.nn.init.kaiming_normal_(m.weight, a=self.hyp['lrelu_g'], nonlinearity='leaky_relu') if not m.bias is None: torch.nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): torch.nn.init.kaiming_normal_(m.weight) if not m.bias is None: torch.nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): torch.nn.init.constant_(m.weight, 1) torch.nn.init.constant_(m.bias, 0)
def train(config_path, cuda): """ Training DeepLab by v2 protocol """ # Configuration CONFIG = Dict(yaml.load(config_path)) device = get_device(cuda) torch.backends.cudnn.benchmark = True # Dataset dataset = get_dataset(CONFIG.DATASET.NAME)( root=CONFIG.DATASET.ROOT, split=CONFIG.DATASET.SPLIT.TRAIN, ignore_label=CONFIG.DATASET.IGNORE_LABEL, mean_bgr=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R), augment=True, base_size=CONFIG.IMAGE.SIZE.BASE, crop_size=CONFIG.IMAGE.SIZE.TRAIN, scales=CONFIG.DATASET.SCALES, flip=True, ) print(dataset) # DataLoader loader = torch.utils.data.DataLoader( dataset=dataset, batch_size=CONFIG.SOLVER.BATCH_SIZE.TRAIN, num_workers=CONFIG.DATALOADER.NUM_WORKERS, shuffle=True, ) loader_iter = iter(loader) # Model check print("Model:", CONFIG.MODEL.NAME) assert ( CONFIG.MODEL.NAME == "DeepLabV2_ResNet101_MSC" ), 'Currently support only "DeepLabV2_ResNet101_MSC"' # Model setup model = eval(CONFIG.MODEL.NAME)(n_classes=CONFIG.DATASET.N_CLASSES) state_dict = torch.load(CONFIG.MODEL.INIT_MODEL) print(" Init:", CONFIG.MODEL.INIT_MODEL) for m in model.base.state_dict().keys(): if m not in state_dict.keys(): print(" Skip init:", m) model.base.load_state_dict(state_dict, strict=False) # to skip ASPP model = nn.DataParallel(model) model.to(device) # Loss definition criterion = nn.CrossEntropyLoss(ignore_index=CONFIG.DATASET.IGNORE_LABEL) criterion.to(device) # Optimizer optimizer = torch.optim.SGD( # cf lr_mult and decay_mult in train.prototxt params=[ { "params": get_params(model.module, key="1x"), "lr": CONFIG.SOLVER.LR, "weight_decay": CONFIG.SOLVER.WEIGHT_DECAY, }, { "params": get_params(model.module, key="10x"), "lr": 10 * CONFIG.SOLVER.LR, "weight_decay": CONFIG.SOLVER.WEIGHT_DECAY, }, { "params": get_params(model.module, key="20x"), "lr": 20 * CONFIG.SOLVER.LR, "weight_decay": 0.0, }, ], momentum=CONFIG.SOLVER.MOMENTUM, ) # Learning rate scheduler scheduler = PolynomialLR( optimizer=optimizer, step_size=CONFIG.SOLVER.LR_DECAY, iter_max=CONFIG.SOLVER.ITER_MAX, power=CONFIG.SOLVER.POLY_POWER, ) # Setup loss logger writer = SummaryWriter(os.path.join(CONFIG.EXP.OUTPUT_DIR, "logs", CONFIG.EXP.ID)) average_loss = MovingAverageValueMeter(CONFIG.SOLVER.AVERAGE_LOSS) # Path to save models checkpoint_dir = os.path.join( CONFIG.EXP.OUTPUT_DIR, "models", CONFIG.EXP.ID, CONFIG.MODEL.NAME.lower(), CONFIG.DATASET.SPLIT.TRAIN, ) makedirs(checkpoint_dir) print("Checkpoint dst:", checkpoint_dir) # Freeze the batch norm pre-trained on COCO model.train() model.module.base.freeze_bn() for iteration in tqdm( range(1, CONFIG.SOLVER.ITER_MAX + 1), total=CONFIG.SOLVER.ITER_MAX, dynamic_ncols=True, ): # Clear gradients (ready to accumulate) optimizer.zero_grad() loss = 0 for _ in range(CONFIG.SOLVER.ITER_SIZE): try: _, images, labels = next(loader_iter) except: loader_iter = iter(loader) _, images, labels = next(loader_iter) # Propagate forward logits = model(images.to(device)) # Loss iter_loss = 0 for logit in logits: # Resize labels for {100%, 75%, 50%, Max} logits _, _, H, W = logit.shape labels_ = resize_labels(labels, size=(H, W)) iter_loss += criterion(logit, labels_.to(device)) # Propagate backward (just compute gradients wrt the loss) iter_loss /= CONFIG.SOLVER.ITER_SIZE iter_loss.backward() loss += float(iter_loss) #print(loss) average_loss.add(loss) # Update weights with accumulated gradients optimizer.step() # Update learning rate scheduler.step(epoch=iteration) # TensorBoard if iteration % CONFIG.SOLVER.ITER_TB == 0: writer.add_scalar("loss/train", average_loss.value()[0], iteration) for i, o in enumerate(optimizer.param_groups): writer.add_scalar("lr/group_{}".format(i), o["lr"], iteration) for i in range(torch.cuda.device_count()): writer.add_scalar( "gpu/device_{}/memory_cached".format(i), torch.cuda.memory_cached(i) / 1024 ** 3, iteration, ) if False: for name, param in model.module.base.named_parameters(): name = name.replace(".", "/") # Weight/gradient distribution writer.add_histogram(name, param, iteration, bins="auto") if param.requires_grad: writer.add_histogram( name + "/grad", param.grad, iteration, bins="auto" ) # Save a model if iteration % CONFIG.SOLVER.ITER_SAVE == 0: torch.save( model.module.state_dict(), os.path.join(checkpoint_dir, "checkpoint_{}.pth".format(iteration)), ) torch.save( model.module.state_dict(), os.path.join(checkpoint_dir, "checkpoint_final.pth") )
def main(config, cuda, gpu): # Configuration CONFIG = Dict(yaml.load(open(config))) # CUDA check cuda = cuda and torch.cuda.is_available() if cuda: gpu_ids = [int(string) for string in gpu.split(',')] current_device = torch.cuda.current_device() print('Running on', torch.cuda.get_device_name(current_device), gpu_ids) # Dataset dataset = CocoStuff10k( root=CONFIG.ROOT, split='train', image_size=513, crop_size=CONFIG.IMAGE.SIZE.TRAIN, scale=True, flip=True, ) # DataLoader loader = torch.utils.data.DataLoader( dataset=dataset, batch_size=CONFIG.BATCH_SIZE, num_workers=CONFIG.NUM_WORKERS, shuffle=True, ) loader_iter = iter(loader) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES) state_dict = torch.load(CONFIG.INIT_MODEL) model.load_state_dict(state_dict, strict=False) # Skip "aspp" layer model = nn.DataParallel(model, device_ids=gpu_ids) if cuda: model.cuda() # Optimizer optimizer = { 'sgd': torch.optim.SGD( # cf lr_mult and decay_mult in train.prototxt params=[{ 'params': get_lr_params(model.module, key='1x'), 'lr': CONFIG.LR, 'weight_decay': CONFIG.WEIGHT_DECAY }, { 'params': get_lr_params(model.module, key='10x'), 'lr': 10 * CONFIG.LR, 'weight_decay': CONFIG.WEIGHT_DECAY }, { 'params': get_lr_params(model.module, key='20x'), 'lr': 20 * CONFIG.LR, 'weight_decay': 0.0 }], momentum=CONFIG.MOMENTUM, ), }.get(CONFIG.OPTIMIZER) # Loss definition criterion = CrossEntropyLoss2d(ignore_index=CONFIG.IGNORE_LABEL) if cuda: criterion.cuda() # TensorBoard Logger writer = SummaryWriter(CONFIG.LOG_DIR) loss_meter = MovingAverageValueMeter(20) model.train() model.module.scale.freeze_bn() for iteration in tqdm( range(1, CONFIG.ITER_MAX + 1), total=CONFIG.ITER_MAX, leave=False, dynamic_ncols=True, ): # Set a learning rate poly_lr_scheduler( optimizer=optimizer, init_lr=CONFIG.LR, iter=iteration - 1, lr_decay_iter=CONFIG.LR_DECAY, max_iter=CONFIG.ITER_MAX, power=CONFIG.POLY_POWER, ) # Clear gradients (ready to accumulate) optimizer.zero_grad() iter_loss = 0 for i in range(1, CONFIG.ITER_SIZE + 1): data, target = next(loader_iter) # Image data = data.cuda() if cuda else data data = Variable(data) # Propagate forward outputs = model(data) # Loss loss = 0 for output in outputs: # Resize target for {100%, 75%, 50%, Max} outputs target_ = resize_target(target, output.size(2)) target_ = target_.cuda() if cuda else target_ target_ = Variable(target_) # Compute crossentropy loss loss += criterion(output, target_) # Backpropagate (just compute gradients wrt the loss) loss /= float(CONFIG.ITER_SIZE) loss.backward() iter_loss += loss.data[0] # Reload dataloader if ((iteration - 1) * CONFIG.ITER_SIZE + i) % len(loader) == 0: loader_iter = iter(loader) loss_meter.add(iter_loss) # Update weights with accumulated gradients optimizer.step() # TensorBoard if iteration % CONFIG.ITER_TF == 0: writer.add_scalar('train_loss', loss_meter.value()[0], iteration) for i, o in enumerate(optimizer.param_groups): writer.add_scalar('train_lr_group{}'.format(i), o['lr'], iteration) if iteration % 1000 != 0: continue for name, param in model.named_parameters(): name = name.replace('.', '/') writer.add_histogram(name, param, iteration, bins="auto") if param.requires_grad: writer.add_histogram(name + '/grad', param.grad, iteration, bins="auto") # Save a model if iteration % CONFIG.ITER_SNAP == 0: torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, 'checkpoint_{}.pth'.format(iteration)), ) # Save a model if iteration % 100 == 0: torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, 'checkpoint_current.pth'), ) torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, 'checkpoint_final.pth'), )
def train(config, cuda): # Auto-tune cuDNN torch.backends.cudnn.benchmark = True # Configuration device = get_device(cuda) CONFIG = Dict(yaml.load(open(config))) # Dataset 10k or 164k dataset = get_dataset(CONFIG.DATASET.NAME)( root=CONFIG.DATASET.ROOT, split=CONFIG.DATASET.SPLIT.TRAIN, base_size=CONFIG.IMAGE.SIZE.TRAIN.BASE, crop_size=CONFIG.IMAGE.SIZE.TRAIN.CROP, mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R), warp=CONFIG.DATASET.WARP_IMAGE, scale=CONFIG.DATASET.SCALES, flip=True, ) # DataLoader loader = torch.utils.data.DataLoader( dataset=dataset, batch_size=CONFIG.SOLVER.BATCH_SIZE.TRAIN, num_workers=CONFIG.DATALOADER.NUM_WORKERS, shuffle=True, ) loader_iter = iter(loader) # Model model = setup_model(CONFIG.MODEL.INIT_MODEL, CONFIG.DATASET.N_CLASSES, train=True) model.to(device) # Optimizer optimizer = torch.optim.SGD( # cf lr_mult and decay_mult in train.prototxt params=[ { "params": get_params(model.module, key="1x"), "lr": CONFIG.SOLVER.LR, "weight_decay": CONFIG.SOLVER.WEIGHT_DECAY, }, { "params": get_params(model.module, key="10x"), "lr": 10 * CONFIG.SOLVER.LR, "weight_decay": CONFIG.SOLVER.WEIGHT_DECAY, }, { "params": get_params(model.module, key="20x"), "lr": 20 * CONFIG.SOLVER.LR, "weight_decay": 0.0, }, ], momentum=CONFIG.SOLVER.MOMENTUM, ) # Learning rate scheduler scheduler = PolynomialLR( optimizer=optimizer, step_size=CONFIG.SOLVER.LR_DECAY, iter_max=CONFIG.SOLVER.ITER_MAX, power=CONFIG.SOLVER.POLY_POWER, ) # Loss definition criterion = nn.CrossEntropyLoss(ignore_index=CONFIG.DATASET.IGNORE_LABEL) criterion.to(device) # TensorBoard logger writer = SummaryWriter(CONFIG.SOLVER.LOG_DIR) average_loss = MovingAverageValueMeter(CONFIG.SOLVER.AVERAGE_LOSS) # Freeze the batch norm pre-trained on COCO model.train() model.module.base.freeze_bn() for iteration in tqdm( range(1, CONFIG.SOLVER.ITER_MAX + 1), total=CONFIG.SOLVER.ITER_MAX, leave=False, dynamic_ncols=True, ): # Clear gradients (ready to accumulate) optimizer.zero_grad() loss = 0 for _ in range(CONFIG.SOLVER.ITER_SIZE): try: images, labels = next(loader_iter) except: loader_iter = iter(loader) images, labels = next(loader_iter) images = images.to(device) labels = labels.to(device) # Propagate forward logits = model(images) # Loss iter_loss = 0 for logit in logits: # Resize labels for {100%, 75%, 50%, Max} logits _, _, H, W = logit.shape labels_ = resize_labels(labels, shape=(H, W)) iter_loss += criterion(logit, labels_) # Backpropagate (just compute gradients wrt the loss) iter_loss /= CONFIG.SOLVER.ITER_SIZE iter_loss.backward() loss += float(iter_loss) average_loss.add(loss) # Update weights with accumulated gradients optimizer.step() # Update learning rate scheduler.step(epoch=iteration) # TensorBoard if iteration % CONFIG.SOLVER.ITER_TB == 0: writer.add_scalar("loss/train", average_loss.value()[0], iteration) for i, o in enumerate(optimizer.param_groups): writer.add_scalar("lr/group{}".format(i), o["lr"], iteration) if False: # This produces a large log file for name, param in model.named_parameters(): name = name.replace(".", "/") # Weight/gradient distribution writer.add_histogram(name, param, iteration, bins="auto") if param.requires_grad: writer.add_histogram(name + "/grad", param.grad, iteration, bins="auto") # Save a model if iteration % CONFIG.SOLVER.ITER_SAVE == 0: torch.save( model.module.state_dict(), osp.join(CONFIG.MODEL.SAVE_DIR, "checkpoint_{}.pth".format(iteration)), ) # To verify progress separately torch.save( model.module.state_dict(), osp.join(CONFIG.MODEL.SAVE_DIR, "checkpoint_current.pth"), ) torch.save( model.module.state_dict(), osp.join(CONFIG.MODEL.SAVE_DIR, "checkpoint_final.pth"), )
def train(**kwargs): # first free all GPU memory t.cuda.empty_cache() """ Get options """ opt = Config() print_options(opt) # overwrite options from commandline for k_, v_ in kwargs.items(): setattr(opt, k_, v_) device = t.device('cuda') if opt.gpu else t.device('cpu') # TODO: visualization """ Dataset """ dataset = create_dataset(opt) dataset_size = len(dataset) iter_per_epoch = int(dataset_size / opt.batch_size) print(f'loaded {dataset_size} images for training') """ Create Network Instances """ model_names = ['netG_x', 'netG_y', 'netD_x', 'netD_y'] netG_x = ResnetGenerator(opt) netG_y = ResnetGenerator(opt) # print(netG_x) netD_x = NLayerDiscriminator(opt) netD_y = NLayerDiscriminator(opt) # print(netD_x) if opt.gpu: netG_x.to(device) summary(netG_x, input_size=(3, opt.crop_size, opt.crop_size)) netG_y.to(device) netD_x.to(device) summary(netD_x, input_size=(3, opt.crop_size, opt.crop_size)) netD_y.to(device) """ Define optimizer and Loss """ optimizer_g = t.optim.Adam(itertools.chain(netG_x.parameters(), netG_y.parameters()), lr=opt.g_lr, betas=(opt.beta1, 0.999)) optimizer_d = t.optim.Adam(itertools.chain(netD_x.parameters(), netD_y.parameters()), lr=opt.d_lr, betas=(opt.beta1, 0.999)) optimizers = [optimizer_g, optimizer_d] """ Forward cycle loss: lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper) Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper) Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper) """ lambda_X = 10.0 # weight for cycle loss (A -> B -> A^) lambda_Y = 10.0 # weight for cycle loss (B -> A -> B^) lambda_identity = 0.5 # 定义 GAN 损失,define GAN loss. # it's a MSELoss() when initialized, only calculate later during iteration # criterionGAN = nn.MSELoss().to(device) criterionGAN = GANLoss(gan_mode='lsgan') # cycle loss criterionCycle = nn.L1Loss() # identical loss criterionIdt = nn.L1Loss() # loss meters loss_X_meter = MovingAverageValueMeter(opt.plot_every) loss_Y_meter = MovingAverageValueMeter(opt.plot_every) score_Dx_real_y = MovingAverageValueMeter(opt.plot_every) score_Dx_fake_y = MovingAverageValueMeter(opt.plot_every) losses = {} scores = {} """ use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1 """ for epoch in range(opt.max_epochs): epoch_start_time = time.time() """ calculate losses, gradients, and update network weights; called in every iteration """ for i, data in enumerate(dataset): real_x = data['A'].to(device) real_y = data['B'].to(device) ###################### # X -> Y' -> X^ cycle ###################### optimizer_g.zero_grad() # set g_x and g_y gradients to zero fake_y = netG_x(real_x) # X -> Y' prediction = netD_x(fake_y) #netD_x provide feedback to netG_x loss_G_X = criterionGAN(prediction, True) # cycle_consistance x_hat = netG_y(fake_y) # Y' -> X^ # Forward cycle loss x^ = || G_y(G_x(real_x)) || loss_cycle_X = criterionCycle(x_hat, real_x) * lambda_X # identity loss if lambda_identity > 0: # netG_x should be identity if real_y is fed: ||netG_x(real_y) - real_y|| idt_x = netG_x(real_y) loss_idt_x = criterionIdt(idt_x, real_y) * lambda_Y * lambda_identity else: loss_idt_x = 0. loss_X = loss_G_X + loss_cycle_X + loss_idt_x loss_X.backward(retain_graph=True) optimizer_g.step() loss_X_meter.add(loss_X.item()) ###################### # Y -> X' -> Y^ cycle ###################### optimizer_g.zero_grad() # set g_x and g_y gradients to zero fake_x = netG_y(real_y) # Y -> X' prediction = netD_y(fake_x) loss_G_Y = criterionGAN(prediction, True) # print(f'loss_G_Y = {round(float(loss_G_Y), 3)}') y_hat = netG_x(fake_x) # Y -> X' -> Y^ # Forward cycle loss y^ = || G_x(G_y(real_y)) || loss_cycle_Y = criterionCycle(y_hat, real_y) * lambda_Y # identity loss if lambda_identity > 0: # netG_y should be identiy if real_x is fed: ||netG_y(real_x) - real_x|| idt_y = netG_y(real_x) loss_idt_y = criterionIdt(idt_y, real_x) * lambda_X * lambda_identity else: loss_idt_y = 0. loss_Y = loss_G_Y + loss_cycle_Y + loss_idt_y loss_Y.backward(retain_graph=True) optimizer_g.step() loss_Y_meter.add(loss_Y.item()) ###################### # netD_x ###################### optimizer_d.zero_grad() # loss_real pred_real = netD_x(real_y) loss_D_x_real = criterionGAN(pred_real, True) score_Dx_real_y.add(float(pred_real.data.mean())) # loss_fake pred_fake = netD_x(fake_y) loss_D_x_fake = criterionGAN(pred_fake, False) score_Dx_fake_y.add(float(pred_fake.data.mean())) # loss and backward loss_D_x = (loss_D_x_real + loss_D_x_fake) * 0.5 loss_D_x.backward() optimizer_d.step() ###################### # netD_y ###################### optimizer_d.zero_grad() # loss_real pred_real = netD_y(real_x) loss_D_y_real = criterionGAN(pred_real, True) # loss_fake pred_fake = netD_y(fake_x) loss_D_y_fake = criterionGAN(pred_fake, False) # loss and backward loss_D_y = (loss_D_y_real + loss_D_y_fake) * 0.5 loss_D_y.backward() optimizer_d.step() # save snapshot if i % opt.plot_every == 0: filename = opt.name + '_snap_%03d_%05d.png' % ( epoch, i, ) test_path = os.path.join(opt.checkpoint_path, filename) tv.utils.save_image(fake_y, test_path, normalize=True) print(f'{filename} saved.') losses['loss_X'] = loss_X_meter.value()[0] losses['loss_Y'] = loss_Y_meter.value()[0] scores['score_Dx_real_y'] = score_Dx_real_y.value()[0] scores['score_Dx_fake_y'] = score_Dx_fake_y.value()[0] print(losses) print(scores) # print(f'iteration {i} finished') # save model if epoch % opt.save_every == 0 or epoch == opt.max_epochs - 1: save_filename = f'{opt.name}_netG_{epoch}.pth' save_filepath = os.path.join(opt.model_path, save_filename) t.save(netG_x.state_dict(), save_filepath) print(f'model saved as {save_filename}') # epoch end logs epoech_time = int(time.time() - epoch_start_time) print_options(opt, epoch_log=True, epoch=epoch, time=epoech_time, losses=losses, scores=scores) print()
def train(**kwargs): opt._parse(kwargs) image_folder_path = 'DataSets/images/' cvs_file_path = 'DataSets/labels.csv' dataset = DataSets(cvs_file_path, image_folder_path) data_size = len(dataset) indices = list(range(data_size)) split = int(np.floor(data_size * 0.2)) np.random.seed(42) np.random.shuffle(indices) train_indices, val_indices = indices[split:], indices[:split] train_sampler = torch.utils.data.SubsetRandomSampler(train_indices) valid_sampler = torch.utils.data.SubsetRandomSampler(val_indices) train_loader = torch.utils.data.DataLoader(dataset, batch_size=1, sampler=train_sampler) val_loader = torch.utils.data.DataLoader(dataset, batch_size=1, sampler=valid_sampler) print('load data') avg_loss = AverageValueMeter() ma20_loss = MovingAverageValueMeter(windowsize=20) faster_rcnn = FasterRCNNVGG16() print('model construct completed') start_epoch = 0 best_map = -100 trainer = FasterRCNNTrainer(faster_rcnn).cuda() optimizer = optim.SGD(trainer.faster_rcnn.parameters(), lr=opt.lr, momentum=0.9) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5) if opt.load_path: print('load pretrained model from %s' % opt.load_path) checkpoint = torch.load(opt.load_path) start_epoch = checkpoint['epoch'] best_map = checkpoint['best_map'] trainer.faster_rcnn.load_state_dict(checkpoint['model_state']) optimizer.load_state_dict(checkpoint['optimizer_state']) print("> Loaded checkpoint '{}' (epoch {})".format( args.resume, start_epoch)) #trainer.vis.text(dataset.db.label_names, win='labels') # set tensor-board for visualization writer = SummaryWriter('runs/' + opt.log_root) for epoch in range(start_epoch, opt.epoch): trainer.train(mode=True) #must set as that in tranning for ii, (img, _, _, bbox_, label_, scale, _) in enumerate(train_loader): scale = at.scalar(scale) img, bbox, label = img.cuda().float(), bbox_.cuda(), label_.cuda() optimizer.zero_grad() loss = trainer.forward(img, bbox, label, scale) loss.total_loss.backward() optimizer.step() #print(loss) #print(loss.total_loss) loss_value = loss.total_loss.cpu().data.numpy() avg_loss.add(float(loss_value)) ma20_loss.add(float(loss_value)) print( '[epoch:{}/{}] [batch:{}/{}] [sample_loss:{:.4f}] [avg_loss:{:.4f}] [ma20_loss:{:.4f}]' .format(epoch, opt.epoch, ii + 1, len(train_loader), loss.total_loss.data, avg_loss.value()[0], ma20_loss.value()[0])) if (ii + 1) % opt.plot_every == 0: niter = epoch * len(train_loader) + ii writer.add_scalar('Train/Loss', ma20_loss.value()[0], niter) eval_result = eval(val_loader, faster_rcnn, test_num=opt.test_num) print(eval_result['map']) if eval_result['map'] > best_map: best_map = eval_result['map'] state = { "epoch": epoch + 1, "best_map": best_map, "model_state": trainer.faster_rcnn.state_dict(), "optimizer_state": optimizer.state_dict() } torch.save(state, opt.model_para) scheduler.step() state = { "epoch": epoch + 1, "best_map": best_map, "model_state": trainer.faster_rcnn.state_dict(), "optimizer_state": optimizer.state_dict() } torch.save(state, 'last_epoch.pkl') writer.close()
def main(): global global_iter_idx args = parse_args() if args.verbose: default_handler = logging.StreamHandler(sys.stdout) logger.addHandler(default_handler) logger.setLevel(logging.DEBUG) if args.log_file is not None: logfile_handler = logging.FileHandler(args.log_file) logger_formatter = logging.Formatter( '%(name)s - %(levelname)s - %(message)s') logfile_handler.setFormatter(logger_formatter) logger.addHandler(logfile_handler) logger.setLevel(logging.DEBUG) train_data_source = davis.DavisDataset(base_dir=os.path.join( root_dir, 'dataset', 'DAVIS'), image_size=args.image_dims, year=2016, phase='train', transform=davis.ToTensor()) train_triplet_sampler = davis.TripletSampler(dataset=train_data_source, num_triplets=args.batch_size, randomize=True) train_data_loader = DataLoader(dataset=train_data_source, batch_sampler=train_triplet_sampler) val_data_source = davis.DavisDataset(base_dir=os.path.join( root_dir, 'dataset', 'DAVIS'), image_size=args.image_dims, year=2016, phase='val', transform=davis.ToTensor()) val_triplet_sampler = davis.TripletSampler( dataset=val_data_source, num_triplets=args.num_val_batches, randomize=True) val_data_loader = DataLoader(dataset=val_data_source, batch_sampler=val_triplet_sampler) model = network.BFVOSNet(embedding_vector_dims=args.embedding_vector_dims) train_loss_fn = loss.MinTripletLoss(alpha=args.alpha) val_loss_fn = loss.validation_loss if has_cuda: model = model.cuda() train_loss_fn = train_loss_fn.cuda() train_loss_fn.to(device) logger.debug("Model and loss function moved to CUDA") start_epoch = 0 if args.checkpoint_path is not None: epoch_substr = args.checkpoint_path.split('epoch_')[1] start_epoch = int(epoch_substr.split('_')[0]) batch_substr = epoch_substr.split('_')[1].split('batch_') if len(batch_substr) > 1: global_iter_idx = int(batch_substr[1].split('_')[0]) if has_cuda: if args.checkpoint_path is not None: # Load pre-trained weights for entire model model.load_state_dict( torch.load( args.checkpoint_path, map_location=lambda storage, loc: storage.cuda(gpu_id))) logger.info("Loaded checkpoint from {}".format( args.checkpoint_path)) else: # Load pre-trained weights for feature extraction head model.load_state_dict(torch.load( deeplab_resnet_pre_trained_path, map_location=lambda storage, loc: storage.cuda(gpu_id)), strict=False) logger.info("Loaded DeepLab ResNet from {}".format( deeplab_resnet_pre_trained_path)) else: if args.checkpoint_path is not None: model.load_state_dict(torch.load(args.checkpoint_path)) logger.info("Loaded checkpoint from {}".format( args.checkpoint_path)) else: model.load_state_dict(torch.load(deeplab_resnet_pre_trained_path), strict=False) logger.info("Loaded DeepLab ResNet from {}".format( deeplab_resnet_pre_trained_path)) # Load to appropriate device and set to training mode but freeze feature extraction layer model.to(device).train() model.freeze_feature_extraction() # Initialize optimizer to train only the unfrozen layers optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=args.learning_rate, momentum=args.momentum) # Initialize meter and writer train_loss_meter = MovingAverageValueMeter(20) val_loss_meter = AverageValueMeter() summary_writer = SummaryWriter(tensorboard_save_dir) # Train for epoch in tqdm(range(start_epoch, args.num_epochs)): logger.info("Epoch {}/{}".format(epoch + 1, args.num_epochs)) train(epoch, train_data_loader, val_data_loader, model, train_loss_fn, val_loss_fn, optimizer, train_loss_meter, val_loss_meter, summary_writer, args.log_interval, args.checkpoint_interval, args.val_interval, args.num_val_batches) # Save final model after all epochs model.eval().cpu() save_model_filename = "epoch_{}_{}.model".format( args.num_epochs, str(time.time()).replace(" ", "_").replace(".", "_")) save_model_path = os.path.join(model_dir, save_model_filename) torch.save(model.state_dict(), save_model_path) logger.info("Model saved to {}".format(save_model_filename)) training_config_save_path = os.path.join( config_save_dir, save_model_filename.replace('.model', '.json')) training_config = vars(args) training_config['device'] = str(torch.device) with open(training_config_save_path, 'w') as f: json.dump(training_config, f) logger.info( "Training config saved to {}".format(training_config_save_path))
save_dir, 'models', modelName + '_epoch-' + str(resume_epoch - 1) + '.pth'), map_location=lambda storage, loc: storage) ) # Load all tensors onto the CPU if gpu_id >= 0: torch.cuda.set_device(device=gpu_id) net.cuda() if resume_epoch != nEpochs: # Logging into Tensorboard log_dir = os.path.join( save_dir, 'models', datetime.now().strftime('%b%d_%H-%M-%S') + '_' + socket.gethostname()) writer = SummaryWriter(log_dir=log_dir) loss_meter = MovingAverageValueMeter(20) # Use the following optimizer optimizer = optim.SGD(net.parameters(), lr=p['lr'], momentum=p['momentum'], weight_decay=p['wd']) p['optimizer'] = str(optimizer) composed_transforms_tr = transforms.Compose([ tr.RandomSized(512), tr.RandomRotate(15), tr.RandomHorizontalFlip(), tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), tr.ToTensor() ])
def main(config, cuda, excludeval, embedding, continue_from, nolog, inputmix, imagedataset, experimentid, nshot, ishot): frame = inspect.currentframe() args, _, _, values = inspect.getargvalues(frame) #print(values) #in case you want to save to the location of script you're running datadir = os.path.join( '/home/SharedData/omkar/zscoseg/yash_manas/data/datasets', imagedataset) if not nolog: #name the savedir, might add logs/ before the datetime for clarity if experimentid is None: savedir = time.strftime('%Y%m%d%H%M%S') else: savedir = experimentid #the full savepath is then: savepath = os.path.join('logs', imagedataset, savedir) #in case the folder has not been created yet / except already exists error: try: os.makedirs(savepath) print("Log dir:", savepath) except: pass if continue_from is None: #now join the path in save_screenshot: shutil.copytree('./libs/', savepath + '/libs') shutil.copy2(osp.abspath(inspect.stack()[0][1]), savepath) shutil.copy2(config, savepath) args_dict = {} for a in args: args_dict[a] = values[a] with open(savepath + '/args.json', 'w') as fp: json.dump(args_dict, fp) cuda = cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") if cuda: current_device = torch.cuda.current_device() print("Running on", torch.cuda.get_device_name(current_device)) else: print("Running on CPU") # Configuration CONFIG = Dict(yaml.load(open(config), Loader=yaml.FullLoader)) visibility_mask = {} if excludeval: seen_classes = np.load(datadir + '/split/seen_cls.npy') else: seen_classes = np.asarray(np.concatenate([ np.load(datadir + '/split/seen_cls.npy'), np.load(datadir + '/split/val_cls.npy') ]), dtype=int) novel_classes = np.load(datadir + '/split/novel_cls.npy') seen_novel_classes = np.concatenate([seen_classes, novel_classes]) seen_map = np.array([-1] * 256) for i, n in enumerate(list(seen_classes)): seen_map[n] = i visibility_mask[0] = seen_map.copy() for i, n in enumerate(list(novel_classes)): visibility_mask[i + 1] = seen_map.copy() visibility_mask[i + 1][n] = seen_classes.shape[0] + i if excludeval: train = np.load(datadir + '/split/train_list.npy')[:-CONFIG.VAL_SIZE] else: train = np.load(datadir + '/split/train_list.npy') novelset = [] seenset = [] if inputmix == 'novel' or inputmix == 'both': inverse_dict = pickle.load( open(datadir + '/split/inverse_dict_train.pkl', 'rb')) for icls, key in enumerate(novel_classes): if (inverse_dict[key].size > 0): for v in inverse_dict[key][ishot * 20:ishot * 20 + nshot]: novelset.append((v, icls)) #print((v, icls)) if inputmix == 'both': seenset = [] inverse_dict = pickle.load( open(datadir + '/split/inverse_dict_train.pkl', 'rb')) for icls, key in enumerate(seen_classes): if (inverse_dict[key].size > 0): for v in inverse_dict[key][ishot * 20:ishot * 20 + nshot]: seenset.append(v) if inputmix == 'seen': seenset = range(train.shape[0]) sampler = RandomImageSampler(seenset, novelset) if inputmix == 'novel': visible_classes = seen_novel_classes if nshot is not None: nshot = str(nshot) + 'n' elif inputmix == 'seen': visible_classes = seen_classes if nshot is not None: nshot = str(nshot) + 's' elif inputmix == 'both': visible_classes = seen_novel_classes if nshot is not None: nshot = str(nshot) + 'b' print("Visible classes:", visible_classes.size, " \nClasses are: ", visible_classes, "\nTrain Images:", train.shape[0]) #a Dataset 10k or 164k dataset = get_dataset(CONFIG.DATASET)(train=train, test=None, root=CONFIG.ROOT, split=CONFIG.SPLIT.TRAIN, base_size=513, crop_size=CONFIG.IMAGE.SIZE.TRAIN, mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R), warp=CONFIG.WARP_IMAGE, scale=(0.5, 1.5), flip=True, visibility_mask=visibility_mask) # DataLoader loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=CONFIG.BATCH_SIZE.TRAIN, num_workers=CONFIG.NUM_WORKERS, sampler=sampler) if embedding == 'word2vec': class_emb = pickle.load( open(datadir + '/word_vectors/word2vec.pkl', "rb")) elif embedding == 'fasttext': class_emb = pickle.load( open(datadir + '/word_vectors/fasttext.pkl', "rb")) elif embedding == 'fastnvec': class_emb = np.concatenate([ pickle.load(open(datadir + '/word_vectors/fasttext.pkl', "rb")), pickle.load(open(datadir + '/word_vectors/word2vec.pkl', "rb")) ], axis=1) else: print("invalid emb ", embedding) sys.exit() print((class_emb.shape)) class_emb = F.normalize(torch.tensor(class_emb), p=2, dim=1).cuda() loader_iter = iter(loader) DeepLab = DeepLabV2_ResNet101_MSC #import ipdb; ipdb.set_trace() state_dict = torch.load(CONFIG.INIT_MODEL) # Model load model = DeepLab(class_emb.shape[1], class_emb[visible_classes]) if continue_from is not None and continue_from > 0: print("Loading checkpoint: {}".format(continue_from)) #import ipdb; ipdb.set_trace() model = nn.DataParallel(model) state_file = osp.join(savepath, "checkpoint_{}.pth".format(continue_from)) if osp.isfile(state_file + '.tar'): state_dict = torch.load(state_file + '.tar') model.load_state_dict(state_dict['state_dict'], strict=True) elif osp.isfile(state_file): state_dict = torch.load(state_file) model.load_state_dict(state_dict, strict=True) else: print("Checkpoint {} not found".format(continue_from)) sys.exit() else: model.load_state_dict( state_dict, strict=False ) # make strict=True to debug if checkpoint is loaded correctly or not if performance is low model = nn.DataParallel(model) model.to(device) # Optimizer optimizer = { "sgd": torch.optim.SGD( # cf lr_mult and decay_mult in train.prototxt params=[{ "params": get_params(model.module, key="1x"), "lr": CONFIG.LR, "weight_decay": CONFIG.WEIGHT_DECAY, }, { "params": get_params(model.module, key="10x"), "lr": 10 * CONFIG.LR, "weight_decay": CONFIG.WEIGHT_DECAY, }, { "params": get_params(model.module, key="20x"), "lr": 20 * CONFIG.LR, "weight_decay": 0.0, }], momentum=CONFIG.MOMENTUM, ), "adam": torch.optim.Adam( # cf lr_mult and decay_mult in train.prototxt params=[{ "params": get_params(model.module, key="1x"), "lr": CONFIG.LR, "weight_decay": CONFIG.WEIGHT_DECAY, }, { "params": get_params(model.module, key="10x"), "lr": 10 * CONFIG.LR, "weight_decay": CONFIG.WEIGHT_DECAY, }, { "params": get_params(model.module, key="20x"), "lr": 20 * CONFIG.LR, "weight_decay": 0.0, }]) # Add any other optimizer }.get(CONFIG.OPTIMIZER) if 'optimizer' in state_dict: optimizer.load_state_dict(state_dict['optimizer']) print("Learning rate:", CONFIG.LR) # Loss definition criterion = nn.CrossEntropyLoss(ignore_index=-1) criterion.to(device) if not nolog: # TensorBoard Logger if continue_from is not None: writer = SummaryWriter( savepath + '/runs/fs_{}_{}_{}'.format(continue_from, nshot, ishot)) else: writer = SummaryWriter(savepath + '/runs') loss_meter = MovingAverageValueMeter(20) model.train() model.module.scale.freeze_bn() pbar = tqdm( range(1, CONFIG.ITER_MAX + 1), total=CONFIG.ITER_MAX, leave=False, dynamic_ncols=True, ) for iteration in pbar: # Set a learning rate poly_lr_scheduler( optimizer=optimizer, init_lr=CONFIG.LR, iter=iteration - 1, lr_decay_iter=CONFIG.LR_DECAY, max_iter=CONFIG.ITER_MAX, power=CONFIG.POLY_POWER, ) # Clear gradients (ready to accumulate) optimizer.zero_grad() iter_loss = 0 for i in range(1, CONFIG.ITER_SIZE + 1): try: data, target = next(loader_iter) except: loader_iter = iter(loader) data, target = next(loader_iter) # Image data = data.to(device) # Propagate forward outputs = model(data) # Loss loss = 0 for output in outputs: # Resize target for {100%, 75%, 50%, Max} outputs target_ = resize_target(target, output.size(2)) target_ = torch.tensor(target_).to(device) loss += criterion.forward(output, target_) # Backpropagate (just compute gradients wrt the loss) #print(loss) loss /= float(CONFIG.ITER_SIZE) loss.backward() iter_loss += float(loss) del data, target, outputs #print(iter_loss) pbar.set_postfix(loss="%.3f" % iter_loss) # Update weights with accumulated gradients optimizer.step() if not nolog: loss_meter.add(iter_loss) # TensorBoard if iteration % CONFIG.ITER_TB == 0: writer.add_scalar("train_loss", loss_meter.value()[0], iteration) for i, o in enumerate(optimizer.param_groups): writer.add_scalar("train_lr_group{}".format(i), o["lr"], iteration) if False: # This produces a large log file for name, param in model.named_parameters(): name = name.replace(".", "/") writer.add_histogram(name, param, iteration, bins="auto") if param.requires_grad: writer.add_histogram(name + "/grad", param.grad, iteration, bins="auto") # Save a model if continue_from is not None: if iteration in CONFIG.ITER_SAVE: torch.save( { 'iteration': iteration, 'state_dict': model.state_dict(), }, osp.join( savepath, "checkpoint_{}_{}_{}_{}.pth.tar".format( continue_from, nshot, ishot, iteration)), ) # Save a model (short term) [unnecessary for fewshot] if False and iteration % 100 == 0: torch.save( { 'iteration': iteration, 'state_dict': model.state_dict(), }, osp.join( savepath, "checkpoint_{}_{}_{}_current.pth.tar".format( continue_from, nshot, ishot)), ) print( osp.join( savepath, "checkpoint_{}_{}_{}_current.pth.tar".format( continue_from, nshot, ishot))) else: if iteration % CONFIG.ITER_SAVE == 0: torch.save( { 'iteration': iteration, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), }, osp.join(savepath, "checkpoint_{}.pth.tar".format(iteration)), ) # Save a model (short term) if iteration % 100 == 0: torch.save( { 'iteration': iteration, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), }, osp.join(savepath, "checkpoint_current.pth.tar"), ) torch.cuda.empty_cache() if not nolog: if continue_from is not None: torch.save( { 'iteration': iteration, 'state_dict': model.state_dict(), }, osp.join( savepath, "checkpoint_{}_{}_{}_{}.pth.tar".format( continue_from, nshot, ishot, CONFIG.ITER_MAX))) else: torch.save( { 'iteration': iteration, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), }, osp.join(savepath, "checkpoint_{}.pth.tar".format(CONFIG.ITER_MAX)))
def main(): parser = argparse.ArgumentParser(description="Train the SharpNet network") parser.add_argument('-c', '--configFile', required=True, help='Path to config yaml file', metavar='path/to/config') args = parser.parse_args() CONFIG_FILE_PATH = args.configFile with open(CONFIG_FILE_PATH) as fd: config_yaml = oyaml.load( fd) # Returns an ordered dict. Used for printing config = AttrDict(config_yaml) print( colored( 'Config being used for training:\n{}\n\n'.format( oyaml.dump(config_yaml)), 'green')) os.environ['CUDA_VISIBLE_DEVICES'] = config.train.cuda_device cuda = False if config.train.nocuda else True resnet50_url = 'https://download.pytorch.org/models/resnet50-19c8e357.pth' cuda = cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") if cuda: current_device = torch.cuda.current_device() print("Running on " + torch.cuda.get_device_name(current_device)) else: print("Running on CPU") now = datetime.datetime.now() date_str = now.strftime("%d-%m-%Y_%H-%M") t = [] torch.manual_seed(329) bias = True if config.train.bias else False # build model model = SharpNet(ResBlock, [3, 4, 6, 3], [2, 2, 2, 2, 2], use_normals=True if config.train.normals else False, use_depth=True if config.train.depth else False, use_boundary=True if config.train.boundary else False, bias_decoder=bias) model_dict = model.state_dict() # Load pretrained weights resnet_path = 'models/resnet50-19c8e357.pth' if not os.path.exists(resnet_path): command = 'wget ' + resnet50_url + ' && mkdir models/ && mv resnet50-19c8e357.pth models/' os.system(command) resnet50_dict = torch.load(resnet_path) resnet_dict = { k.replace('.', '_img.', 1): v for k, v in resnet50_dict.items() if k.replace('.', '_img.', 1) in model_dict } # load weights up to pool print('Loading checkpoint from {}'.format(config.train.pretrained_model)) if config.train.pretrained_model is not None: model_path = config.train.pretrained_model tmp_dict = torch.load(model_path) if config.train.depth: pretrained_dict = { k: v for k, v in tmp_dict.items() if k in model_dict } else: pretrained_dict = { k: v for k, v in tmp_dict.items() if (k in model_dict and not k.startswith('depth_decoder')) } else: pretrained_dict = resnet_dict try: model_dict.update(pretrained_dict) model.load_state_dict(model_dict) print('Successfully loaded pretrained ResNet weights') except: print('Could not load the pretrained model weights') sys.exit(0) model = nn.DataParallel(model) model.to(device) model.zero_grad() model.train() freeze_decoders = config.train.decoder_freeze.split(',') freeze_model_decoders(model.module, freeze_decoders) if config.train.dataset != 'NYU': sharpnet_loss = SharpNetLoss( lamb=0.5, mu=1.0, use_depth=True if config.train.depth else False, use_boundary=True if config.train.boundary else False, use_normals=True if config.train.normals else False, use_geo_consensus=True if config.train.geo_consensus else False) else: sharpnet_loss = SharpNetLoss( lamb=0.5, mu=1.0, use_depth=True if config.train.depth else False, use_boundary=False, use_normals=False, use_geo_consensus=True if config.train.geo_consensus else False) if config.train.optimizer == 'SGD': optimizer = SGD(params=get_params(model), lr=float(config.train.learning_rate), weight_decay=float(config.train.decay), momentum=0.9) elif config.train.optimizer == 'Adam': optimizer = Adam(params=get_params(model), lr=float(config.train.learning_rate), weight_decay=float(config.train.decay)) else: print( 'Could not configure the optimizer, please select --optimizer Adam or SGD' ) sys.exit(0) # TensorBoard Logger train_loss_meter = MovingAverageValueMeter(20) val_loss_meter = MovingAverageValueMeter(3) depth_loss_meter = MovingAverageValueMeter( 3) if config.train.depth else None normals_loss_meter = MovingAverageValueMeter( 3) if config.train.normals and config.train.dataset != 'NYU' else None grad_loss_meter = MovingAverageValueMeter( 3) if config.train.depth else None boundary_loss_meter = MovingAverageValueMeter( 3) if config.train.boundary and config.train.dataset != 'NYU' else None consensus_loss_meter = MovingAverageValueMeter( 3) if config.train.geo_consensus else None exp_name = config.train.experiment_name if config.train.experiment_name is not None else '' print('Experiment Name: {}'.format(exp_name)) log_dir = os.path.join('logs', 'Joint', str(exp_name) + '_' + date_str) cp_dir = os.path.join('logs', 'Joint', str(exp_name) + '_' + date_str) print('Checkpoint Directory: {}'.format(cp_dir)) train_writer = SummaryWriter(os.path.join(log_dir, 'train')) val_writer = SummaryWriter(os.path.join(log_dir, 'val')) if not os.path.exists(cp_dir): os.makedirs(cp_dir) if not os.path.exists(log_dir): os.makedirs(os.path.join(log_dir, 'train')) os.makedirs(os.path.join(log_dir, 'val')) train_dataloader, val_dataloader = get_trainval_splits( config) # SHREK: Added Modification to pass in config. # Either pass in path to config file, or real yaml and pass in the dict of config file. # Config file need only contain the paths to datasets train and val. # For val, we'd like to pass real images dataset. for epoch in range(config.train.max_epoch): if config.train.optimizer == 'SGD': adjust_learning_rate(float(config.train.learning_rate), config.train.lr_mode, float(config.train.gradient_step), config.train.max_epoch, optimizer, epoch) train_epoch(train_dataloader, val_dataloader, model, sharpnet_loss, optimizer, config.train.start_epoch + epoch, train_writer, val_writer, train_loss_meter, val_loss_meter, depth_loss_meter, grad_loss_meter, normals_loss_meter, date_str=date_str, model_save_path=cp_dir, config=config, boundary_loss_meter=boundary_loss_meter, consensus_loss_meter=consensus_loss_meter) # Save a model if epoch % 1 == 0 and epoch > int(0.9 * config.train.max_epoch): torch.save( model.state_dict(), os.path.join( cp_dir, 'checkpoint_{}_final.pth'.format(config.train.start_epoch + epoch)), ) elif epoch % 1 == 0: torch.save( model.state_dict(), os.path.join( cp_dir, 'checkpoint_{}_final.pth'.format(config.train.start_epoch + epoch)), ) torch.save( model.state_dict(), os.path.join( cp_dir, 'checkpoint_{}_final.pth'.format(config.train.start_epoch + config.train.max_epoch)), ) return None
def train(self): torch.cuda.empty_cache() ###################### # Save / Load model ###################### if self.opt.continue_train: try: self.continue_from_latest_checkpoint() except Exception as e: self.logger.error(e) self.opt.continue_train = False self.reset_save() else: self.reset_save() dataset = UnitRiffDataset(self.opt.dataset_name, self.opt.instr_type) dataset_size = len(dataset) self.logger.info( f'Dataset {self.opt.dataset_name} loaded, size {dataset_size}') ###################### # Initiate ###################### criterionGAN = nn.BCEWithLogitsLoss() GLoss_meter = MovingAverageValueMeter(self.opt.plot_every) DLoss_meter = MovingAverageValueMeter(self.opt.plot_every) losses = {} ###################### # Start Training ###################### for epoch in range(self.opt.start_epoch, self.opt.max_epoch): loader = DataLoader(dataset, batch_size=self.opt.batch_size, shuffle=True, num_workers=self.opt.num_threads, drop_last=False) epoch_start_time = time.time() for i, data in enumerate(loader): batch_size = data.size(0) # print(batch_size) real_label = torch.ones(size=[batch_size, 1], device=self.device) fake_label = torch.zeros(size=[batch_size, 1], device=self.device) seed = np.array([ generate_random_seed(1, self.opt.instr_type, pattern=self.opt.chord_type) for _ in range(batch_size) ]) # print(seed.shape) noise = torch.randn(batch_size, self.opt.seed_size, device=self.device) seed = torch.from_numpy(seed).to(device=self.device, dtype=torch.float) fake_data = self.generator(noise, seed, batch_size) D_fake = self.discriminator(fake_data, batch_size) real_data = torch.unsqueeze(data, 1).to(device=self.device, dtype=torch.float) D_real = self.discriminator(real_data, batch_size) # print(D_fake.shape) ###################### # Generator ###################### self.G_optimizer.zero_grad() loss_G = criterionGAN(D_fake, real_label) loss_G.backward(retain_graph=True) self.G_optimizer.step() self.G_optimizer.zero_grad() loss_G = criterionGAN(D_fake, real_label) loss_G.backward(retain_graph=True) self.G_optimizer.step() GLoss_meter.add(loss_G.item()) ###################### # Discriminator ###################### self.D_optimizer.zero_grad() loss_D_real = criterionGAN(D_real, real_label) loss_D_fake = criterionGAN(D_fake, fake_label) loss_D = 0.5 * loss_D_real + 0.5 * loss_D_fake loss_D.backward() self.D_optimizer.step() DLoss_meter.add(loss_D.item()) if epoch % self.opt.save_every == 0 or epoch == self.opt.max_epoch - 1: self.save_model(epoch) losses['loss_G'] = float(GLoss_meter.value()[0]) losses['loss_D'] = float(DLoss_meter.value()[0]) self.G_scheduler.step(epoch) self.D_scheduler.step(epoch) epoch_time = int(time.time() - epoch_start_time) self.logger.info( f'Epoch {epoch} finished, cost time {epoch_time}\n') self.logger.info(str(losses) + '\n\n')
def train(self): torch.cuda.empty_cache() if self.model == 'base': dataset = SteelyDataset(self.genreA, self.genreB, 'train', use_mix=False) dataset_size = len(dataset) else: dataset = SteelyDataset(self.genreA, self.genreB, 'train', use_mix=True) dataset_size = len(dataset) if self.continue_train: self.continue_from_latest_checkpoint() else: self.empty_checkpoints() self.create_save_dirs() iter_num = int(dataset_size / self.batch_size) print(f'loaded {dataset_size} images for training') # optimizers = [optimizer_g, optimizer_d] lambda_A = 10.0 # weight for cycle loss (A -> B -> A^) lambda_B = 10.0 # weight for cycle loss (B -> A -> B^) L1_lambda = 10.0 lambda_identity = 0.5 # it's a MSELoss() when initialized, only calculate later during iteration # criterionGAN = nn.MSELoss().to(device) criterionGAN = GANLoss(gan_mode='vanilla') # cycle loss criterionCycle = nn.L1Loss() # identical loss criterionIdt = nn.L1Loss() GLoss_meter = MovingAverageValueMeter(self.plot_every) DLoss_meter = MovingAverageValueMeter(self.plot_every) # score_DA_real_B = MovingAverageValueMeter(self.plot_every) # score_DA_fake_B = MovingAverageValueMeter(self.plot_every) # loss meters losses = {} scores = {} for epoch in range(self.start_epoch, self.max_epoch): loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True, num_workers=1, drop_last=True) epoch_start_time = time.time() for i, data in enumerate(loader): real_A = torch.unsqueeze(data[:, 0, :, :], 1).to(self.device, dtype=torch.float) real_B = torch.unsqueeze(data[:, 1, :, :], 1).to(self.device, dtype=torch.float) gaussian_noise = torch.abs( torch.normal(mean=torch.zeros(self.data_shape), std=1)).to(self.device, dtype=torch.float) if self.model == 'base': self.GA2B_optimizer.zero_grad( ) # set g_x and g_y gradients to zero fake_B = self.generator_A2B(real_A) # X -> Y' fake_B_copy = copy.copy(fake_B.detach()) DB_fake = self.discriminator_B( fake_B + gaussian_noise) #netD_x provide feedback to netG_x loss_G_A2B = criterionGAN(DB_fake, True) # cycle_consistence cycle_A = self.generator_B2A(fake_B) # Y' -> X^ # Forward cycle loss x^ = || G_y(G_x(real_x)) || loss_cycle_A2B = criterionCycle(cycle_A, real_A) * lambda_A # identity loss if lambda_identity > 0: # netG_x should be identity if real_y is fed: ||netG_x(real_y) - real_y|| idt_A = self.generator_A2B(real_B) loss_idt_A = criterionIdt( idt_A, real_B) * lambda_A * lambda_identity else: loss_idt_A = 0. loss_A = loss_G_A2B + 5. * loss_cycle_A2B loss_A.backward(retain_graph=True) self.GA2B_optimizer.step() ###################### # B -> A' -> B^ cycle ###################### self.GB2A_optimizer.zero_grad( ) # set g_x and g_y gradients to zero fake_A = self.generator_B2A(real_B) # Y -> X' fake_A_copy = copy.copy(fake_A.detach()) DA_fake = self.discriminator_A(fake_A + gaussian_noise) loss_G_B2A = criterionGAN(DA_fake, True) # print(f'loss_G_Y = {round(float(loss_G_Y), 3)}') cycle_B = self.generator_A2B(fake_A) # Y -> X' -> Y^ # Forward cycle loss y^ = || G_x(G_y(real_y)) || loss_cycle_B2A = criterionCycle(cycle_B, real_B) * lambda_B # identity loss if lambda_identity > 0: # netG_y should be identiy if real_x is fed: ||netG_y(real_x) - real_x|| idt_B = self.generator_B2A(real_A) loss_idt_B = criterionIdt( idt_B, real_A) * lambda_A * lambda_identity else: loss_idt_B = 0. loss_B = loss_G_B2A + 5. * loss_cycle_B2A loss_B.backward(retain_graph=True) self.GB2A_optimizer.step() ###################### # sample ###################### if self.use_image_poll: [fake_A_sample, fake_B_sample] = self.pool([fake_A_copy, fake_B_copy]) ###################### # netD_A ###################### # loss_real DA_real = self.discriminator_A(real_A + gaussian_noise) loss_DA_real = criterionGAN(DA_real, True) score_DA_real_B.add(float(DA_real.data.mean())) # loss fake if self.use_image_poll: DA_fake_sample = self.discriminator_A(fake_A_sample + gaussian_noise) loss_DA_fake = criterionGAN(DA_fake_sample, False) score_DA_fake_B.add(float(DA_fake_sample.data.mean())) else: loss_DA_fake = criterionGAN(DA_fake, False) score_DA_fake_B.add(float(DA_fake.data.mean())) # loss and backward self.DA_optimizer.zero_grad() loss_DA = (loss_DA_real + loss_DA_fake) * 0.5 loss_DA.backward() self.DA_optimizer.step() ###################### # netD_B ###################### # loss_real DB_real = self.discriminator_B(real_B + gaussian_noise) loss_DB_real = criterionGAN(DB_real, True) # loss_fake if self.use_image_poll: DB_fake_sample = self.discriminator_B(fake_B_sample + gaussian_noise) loss_DB_fake = criterionGAN(DB_fake_sample, False) else: loss_DB_fake = criterionGAN(DB_fake, False) # loss and backward self.DB_optimizer.zero_grad() loss_DB = (loss_DB_real + loss_DB_fake) * 0.5 loss_DB.backward() self.DB_optimizer.step() else: real_mixed = torch.unsqueeze(data[:, 2, :, :], 1).to(self.device, dtype=torch.float) ###################### # A -> B' -> A^ cycle ###################### self.GA2B_optimizer.zero_grad( ) # set g_x and g_y gradients to zero fake_B = self.generator_A2B(real_A) # X -> Y' fake_B_copy = copy.copy(fake_B.detach()) DB_fake = self.discriminator_B( fake_B + gaussian_noise) # netD_x provide feedback to netG_x ''' to_binary ''' loss_G_A2B = criterionGAN(DB_fake, True) # cycle_consistence cycle_A = self.generator_B2A(fake_B) # Y' -> X^ # Forward cycle loss x^ = || G_y(G_x(real_x)) || loss_cycle_A2B = criterionCycle(cycle_A, real_A) * lambda_A # identity loss if lambda_identity > 0: # netG_x should be identity if real_y is fed: ||netG_x(real_y) - real_y|| idt_A = self.generator_A2B(real_B) loss_idt_A = criterionIdt( idt_A, real_B) * lambda_A * lambda_identity else: loss_idt_A = 0. loss_A = loss_G_A2B + 5. * loss_cycle_A2B loss_A.backward(retain_graph=True) self.GA2B_optimizer.step() loss_A_meter.add(loss_A.item()) ###################### # B -> A' -> B^ cycle ###################### self.GB2A_optimizer.zero_grad( ) # set g_x and g_y gradients to zero fake_A = self.generator_B2A(real_B) # Y -> X' fake_A_copy = copy.copy(fake_A.detach()) DA_fake = self.discriminator_A(fake_A + gaussian_noise) loss_G_B2A = criterionGAN(DA_fake, True) # print(f'loss_G_Y = {round(float(loss_G_Y), 3)}') cycle_B = self.generator_A2B(fake_A) # Y -> X' -> Y^ # Forward cycle loss y^ = || G_x(G_y(real_y)) || loss_cycle_B2A = criterionCycle(cycle_B, real_B) * lambda_B # identity loss if lambda_identity > 0: # netG_y should be identiy if real_x is fed: ||netG_y(real_x) - real_x|| idt_B = self.generator_B2A(real_A) loss_idt_B = criterionIdt( idt_B, real_A) * lambda_A * lambda_identity else: loss_idt_B = 0. loss_B = loss_G_B2A + 5. * loss_cycle_B2A loss_B.backward(retain_graph=True) self.GB2A_optimizer.step() loss_B_meter.add(loss_B.item()) ###################### # sample ###################### if self.use_image_poll: [fake_A_sample, fake_B_sample] = self.pool([fake_A_copy, fake_B_copy]) ###################### # netD_A $ netD_A_all ###################### # loss_real DA_real = self.discriminator_A(real_A + gaussian_noise) loss_DA_real = criterionGAN(DA_real, True) # score_DA_real_B.add(float(DA_real.data.mean())) DA_real_all = self.discriminator_A_all(real_mixed + gaussian_noise) loss_DA_all_real = criterionGAN(DA_real_all, True) # loss fake if self.use_image_poll: DA_fake_sample = self.discriminator_A(fake_A_sample + gaussian_noise) loss_DA_fake = criterionGAN(DA_fake_sample, False) # score_DA_fake_B.add(float(DA_fake_sample.data.mean())) DA_fake_sample_all = self.discriminator_A_all( fake_A_sample + gaussian_noise) loss_DA_all_fake = criterionGAN( DA_fake_sample_all, False) else: loss_DA_fake = criterionGAN(DA_fake, False) # score_DA_fake_B.add(float(DA_fake.data.mean())) DA_fake_all = self.discriminator_A_all(fake_A_copy + gaussian_noise) loss_DA_all_fake = criterionGAN(DA_fake_all, False) # loss and backward self.DA_optimizer.zero_grad() loss_DA = (loss_DA_real + loss_DA_fake) * 0.5 loss_DA.backward() self.DA_optimizer.step() self.DA_all_optimizer.zero_grad() loss_DA_all = (loss_DA_all_real + loss_DA_all_fake) * 0.5 loss_DA_all.backward() self.DA_all_optimizer.step() ###################### # netD_A_all ###################### self.DA_all_optimizer.zero_grad() # loss_real DA_real_all = self.discriminator_A_all(real_mixed + gaussian_noise) loss_DA_all_real = criterionGAN(DA_real_all, True) # loss fake DA_fake_sample_all = self.discriminator_A_all( fake_A_sample + gaussian_noise) loss_DA_all_fake = criterionGAN(DA_fake_sample_all, False) loss_DA_all = (loss_DA_all_real + loss_DA_all_fake) * 0.5 loss_DA_all.backward() self.DA_all_optimizer.step() ###################### # netD_B & netD_B_all ###################### # loss_real DB_real = self.discriminator_B(real_B + gaussian_noise) loss_DB_real = criterionGAN(DB_real, True) DB_real_all = self.discriminator_B_all(real_mixed + gaussian_noise) loss_DB_all_real = criterionGAN(DB_real_all, True) # loss_fake if self.use_image_poll: DB_fake_sample = self.discriminator_B(fake_B_sample + gaussian_noise) loss_DB_fake = criterionGAN(DB_fake_sample, False) DB_fake_sample_all = self.discriminator_B_all( fake_B_sample + gaussian_noise) loss_DB_all_fake = criterionGAN( DB_fake_sample_all, False) else: loss_DB_fake = criterionGAN(DB_fake, False) DB_fake_all = self.discriminator_B_all(fake_B_copy + gaussian_noise) loss_DB_all_fake = criterionGAN(DB_fake_all, False) # loss and backward self.DB_optimizer.zero_grad() loss_DB = (loss_DB_real + loss_DB_fake) * 0.5 + ( loss_DB_all_real + loss_DB_all_fake) * 0.5 loss_DB.backward() self.DB_optimizer.step() self.DB_all_optimizer.zero_grad() loss_DB_all = (loss_DB_all_real + loss_DB_all_fake) * 0.5 loss_DB_all.backward() self.DB_all_optimizer.step() ###################### # netD_all ###################### # loss_D_all = loss_DB_all + loss_DA_all # loss_D_all.backward(retain_graph=True) ''' ###################### # netD_B_all ###################### self.DB_all_optimizer.zero_grad() # loss_real DB_real_all = self.discriminator_B_all(real_mixed + gaussian_noise) loss_DB_all_real = criterionGAN(DB_real_all, True) # loss fake DB_fake_sample_all = self.discriminator_B_all(fake_B_sample + gaussian_noise) loss_DB_all_fake = criterionGAN(DB_fake_sample_all, False) loss_DB_all = (loss_DB_all_real + loss_DB_all_fake) * 0.5 loss_DB_all.backward() self.DB_all_optimizer.step() ''' # save snapshot if i % self.plot_every == 0: file_name = self.name + '_snap_%03d_%05d.png' % ( epoch, i, ) test_path = os.path.join(self.checkpoint_path, file_name) tv.utils.save_image(fake_B, test_path, normalize=True) print(f'{file_name} saved.') losses['loss_A'] = loss_A_meter.value()[0] losses['loss_B'] = loss_B_meter.value()[0] scores['score_DA_real_B'] = score_DA_real_B.value()[0] scores['score_DA_fake_B'] = score_DA_fake_B.value()[0] print(losses) print(scores) print('Epoch {} progress: {:.2%}\n'.format( epoch, i / iter_num)) # save model if epoch % self.save_every == 0 or epoch == self.max_epoch - 1: self.save_model(epoch) print(f'model saved') self.GA2B_scheduler.step(epoch) self.GB2A_scheduler.step(epoch) self.DA_scheduler.step(epoch) self.DB_scheduler.step(epoch) if self.model != 'base': self.DA_all_scheduler.step(epoch) self.DB_all_scheduler.step(epoch) epoch_time = int(time.time() - epoch_start_time) print_options(self.opt, epoch_log=True, epoch=epoch, time=epoch_time, losses=losses, scores=scores)
lr_decay_epoch=10): """Sets the learning rate to the initial LR decayed by lr_decay_factor every lr_decay_epoch epochs""" if epoch % lr_decay_epoch == 0: lr = init_lr * (lr_decay_factor**(epoch // lr_decay_epoch)) print('LR is set to {}'.format(lr)) for param_group in optimizer.param_groups: param_group['lr'] = lr model = faster_rcnn(20, backbone='vgg16') if torch.cuda.is_available(): model = model.cuda() optimizer = model.get_optimizer(is_adam=False) avg_loss = AverageValueMeter() ma20_loss = MovingAverageValueMeter(windowsize=20) model.train() for epoch in range(15): adjust_learning_rate(optimizer, epoch, 0.001, lr_decay_epoch=10) for i in range(len(trainval_dataset)): img, bbox, label = trainval_dataset[ i] #this is to retrive the data from datasets easy method img = img / 255 loss = model.loss(img, bbox, label) optimizer.zero_grad() loss.backward() optimizer.step() loss_value = loss.cpu().data.numpy()
def train(self): torch.cuda.empty_cache() ###################### # Save / Load model ###################### if self.opt.continue_train: try: self.continue_from_latest_checkpoint() except Exception as e: self.logger.error(e) return else: self.reset_save() self.logger.add_file_logger(self.opt.log_path) ###################### # Dataset ###################### dataset = ClassifierDataset(self.opt.genreA, self.opt.genreB, 'train') test_dataset = ClassifierDataset(self.opt.genreA, self.opt.genreB, 'test') dataset_size = len(dataset) iter_num = int(dataset_size / self.opt.batch_size) plot_every = iter_num // 10 self.logger.info( f'Dataset loaded, genreA: {self.opt.genreA}, genreB: {self.opt.genreB}, total size: {dataset_size}.' ) ###################### # Initiate ###################### softmax_criterion = nn.BCELoss() Loss_meter = MovingAverageValueMeter(self.opt.plot_every) losses = {} ###################### # Start Training ###################### test_data = torch.from_numpy(test_dataset.get_data()).to( self.device, dtype=torch.float) gaussian_noise = torch.normal(mean=torch.zeros(test_data.shape), std=self.opt.gaussian_std).to( self.device, dtype=torch.float) # test_data += gaussian_noise real_test_label = torch.from_numpy(test_dataset.get_labels()).view( -1, 2).to(self.device, dtype=torch.float) for epoch in range(self.opt.start_epoch, self.opt.max_epoch): loader = DataLoader(dataset, batch_size=self.opt.batch_size, shuffle=True, num_workers=self.opt.num_threads, drop_last=True) epoch_start_time = time.time() for i, batch in enumerate(loader): data = batch[0].to(self.device, dtype=torch.float) real_label = batch[1].view(self.opt.batch_size, 2).to(self.device, dtype=torch.float) self.classifier_optimizer.zero_grad() estimate_train = self.classifier(data) loss = softmax_criterion(estimate_train, real_label) loss.backward() self.classifier_optimizer.step() Loss_meter.add(loss.item()) # test if i % plot_every == 0: with torch.no_grad(): estimate_test = self.classifier(test_data) estimate_test = nn.functional.softmax(estimate_test, dim=1) test_prediction = torch.argmax(estimate_test, 1).eq( torch.argmax(real_test_label, 1)) test_accuracy = torch.mean( test_prediction.type(torch.float32)).cpu() self.logger.info( 'Epoch {} progress {:.2%}: Loss: {}, Accuracy: {}\n'. format(epoch, i / iter_num, Loss_meter.value()[0], test_accuracy)) if epoch % self.opt.save_every == 0 or epoch == self.opt.max_epoch - 1: self.save_model(epoch) self.classifier_scheduler.step(epoch) epoch_time = int(time.time() - epoch_start_time) self.logger.info( f'Epoch {epoch} finished, cost time {epoch_time}\n')
train_file = "train_loss.txt" text_file = "test_loss.txt" save_file = "loss_curve.pdf" train_f = open(train_file) train_d = train_f.readlines() train_f.close() valid_f = open(text_file) valid_d = valid_f.readlines() valid_f.close() train_iter = [] train_loss = [] i = 0 ma_loss = MovingAverageValueMeter(windowsize=500) for s in train_d: i = i + 1 t = s.strip().split(' ') t_iter = int(t[0]) ma_loss.add(float(t[1])) if i % 500 == 0: train_iter.append(t_iter) train_loss.append(ma_loss.value()[0]) valid_iter = [] valid_loss = [] i = 0 for s in valid_d: i = i + 1 if i >= 0:
def main(): parser = argparse.ArgumentParser(description="Train the SharpNet network") parser.add_argument('--dataset', '-d', dest='dataset', help='Name of the dataset (MLT, NYUv2 or pix3d)') parser.add_argument('--exp_name', dest='experiment_name', help='Custom name of the experiment', type=str, default=None) parser.add_argument('--batch-size', '-b', dest='batch_size', type=int, default=3, help='Batch size') parser.add_argument('--iter-size', dest='iter_size', type=int, default=3, help='Iteration size (for accumulated gradients)') parser.add_argument('--boundary', action='store_true', help='Use boundary decoder') parser.add_argument('--normals', action='store_true', help='Use normals decoder') parser.add_argument('--depth', action='store_true', help='Use depth decoder') parser.add_argument('--consensus', dest='geo_consensus', action='store_true') parser.add_argument('--freeze', dest='decoder_freeze', default='', type=str, help='Decoders to freeze (comma seperated)') parser.add_argument('--verbose', action='store_true', help='Activate to display loss components terms') parser.add_argument('--rootdir', '-r', dest='root_dir', default='', help='Root Directory of the dataset') parser.add_argument( '--nocuda', action="store_true", help='Use flag to use on CPU only (currently not supported)') parser.add_argument('--lr', dest='learning_rate', type=float, default=1e-5, help='Initial learning rate') parser.add_argument('--lr-mode', dest='lr_mode', default='poly', help='Learning rate decay mode') parser.add_argument('--max-epoch', dest='max_epoch', type=int, default=1000, help='MAXITER') parser.add_argument('--step', '-s', dest='gradient_step', default=5e-2, help='gradient step') parser.add_argument('--cuda', dest='cuda_device', default="0", help='CUDA device ID') parser.add_argument('--cpu', dest='num_workers', default=4) parser.add_argument('--pretrained-model', dest='pretrained_model', default=None, help="Choose a model to fine tune") parser.add_argument('--start_epoch', dest='start_epoch', default=0, type=int, help="Starting epoch") parser.add_argument('--bias', action="store_true", help="Flag to learn bias in decoder convnet") parser.add_argument('--optimizer', dest='optimizer', default='SGD', type=str, help="Optimizer type: SGD / Adam") parser.add_argument('--decay', dest='decay', default=5e-5, type=float, help="Weight decay rate") args = parser.parse_args() os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device) cuda = False if args.nocuda else True resnet50_url = 'https://download.pytorch.org/models/resnet50-19c8e357.pth' cuda = cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") if cuda: current_device = torch.cuda.current_device() print("Running on " + torch.cuda.get_device_name(current_device)) else: print("Running on CPU") now = datetime.datetime.now() date_str = now.strftime("%d-%m-%Y_%H-%M") t = [] torch.manual_seed(329) bias = True if args.bias else False # build model model = SharpNet(ResBlock, [3, 4, 6, 3], [2, 2, 2, 2, 2], use_normals=True if args.normals else False, use_depth=True if args.depth else False, use_boundary=True if args.boundary else False, bias_decoder=bias) model_dict = model.state_dict() # Load pretrained weights resnet_path = 'models/resnet50-19c8e357.pth' if not os.path.exists(resnet_path): command = 'wget ' + resnet50_url + ' && mkdir models/ && mv resnet50-19c8e357.pth models/' os.system(command) resnet50_dict = torch.load(resnet_path) resnet_dict = { k.replace('.', '_img.', 1): v for k, v in resnet50_dict.items() if k.replace('.', '_img.', 1) in model_dict } # load weights up to pool if args.pretrained_model is not None: model_path = args.pretrained_model tmp_dict = torch.load(model_path) if args.depth: pretrained_dict = { k: v for k, v in tmp_dict.items() if k in model_dict } else: pretrained_dict = { k: v for k, v in tmp_dict.items() if (k in model_dict and not k.startswith('depth_decoder')) } else: pretrained_dict = resnet_dict try: model_dict.update(pretrained_dict) model.load_state_dict(model_dict) print('Successfully loaded pretrained ResNet weights') except: print('Could not load the pretrained model weights') sys.exit(0) model.to(device) model.zero_grad() model.train() freeze_decoders = args.decoder_freeze.split(',') freeze_model_decoders(model, freeze_decoders) if args.dataset != 'NYU': sharpnet_loss = SharpNetLoss( lamb=0.5, mu=1.0, use_depth=True if args.depth else False, use_boundary=True if args.boundary else False, use_normals=True if args.normals else False, use_geo_consensus=True if args.geo_consensus else False) else: sharpnet_loss = SharpNetLoss( lamb=0.5, mu=1.0, use_depth=True if args.depth else False, use_boundary=False, use_normals=False, use_geo_consensus=True if args.geo_consensus else False) if args.optimizer == 'SGD': optimizer = SGD(params=get_params(model), lr=args.learning_rate, weight_decay=args.decay, momentum=0.9) elif args.optimizer == 'Adam': optimizer = Adam(params=get_params(model), lr=args.learning_rate, weight_decay=args.decay) else: print( 'Could not configure the optimizer, please select --optimizer Adam or SGD' ) sys.exit(0) # TensorBoard Logger train_loss_meter = MovingAverageValueMeter(20) val_loss_meter = MovingAverageValueMeter(3) depth_loss_meter = MovingAverageValueMeter(3) if args.depth else None normals_loss_meter = MovingAverageValueMeter( 3) if args.normals and args.dataset != 'NYU' else None grad_loss_meter = MovingAverageValueMeter(3) if args.depth else None boundary_loss_meter = MovingAverageValueMeter( 3) if args.boundary and args.dataset != 'NYU' else None consensus_loss_meter = MovingAverageValueMeter( 3) if args.geo_consensus else None exp_name = args.experiment_name if args.experiment_name is not None else '' print('Experiment Name: {}'.format(exp_name)) log_dir = os.path.join('logs', 'Joint', str(exp_name) + '_' + date_str) cp_dir = os.path.join('checkpoints', 'Joint', str(exp_name) + '_' + date_str) print('Checkpoint Directory: {}'.format(cp_dir)) train_writer = SummaryWriter(os.path.join(log_dir, 'train')) val_writer = SummaryWriter(os.path.join(log_dir, 'val')) if not os.path.exists(cp_dir): os.makedirs(cp_dir) if not os.path.exists(log_dir): os.makedirs(os.path.join(log_dir, 'train')) os.makedirs(os.path.join(log_dir, 'val')) train_dataloader, val_dataloader = get_trainval_splits(args) for epoch in range(args.max_epoch): if args.optimizer == 'SGD': adjust_learning_rate(args.learning_rate, args.lr_mode, args.gradient_step, args.max_epoch, optimizer, epoch) train_epoch(train_dataloader, val_dataloader, model, sharpnet_loss, optimizer, args.start_epoch + epoch, train_writer, val_writer, train_loss_meter, val_loss_meter, depth_loss_meter, grad_loss_meter, normals_loss_meter, date_str=date_str, model_save_path=cp_dir, args=args, boundary_loss_meter=boundary_loss_meter, consensus_loss_meter=consensus_loss_meter) # Save a model if epoch % 2 == 0 and epoch > int(0.9 * args.max_epoch): torch.save( model.state_dict(), os.path.join( cp_dir, 'checkpoint_{}_final.pth'.format(args.start_epoch + epoch)), ) elif epoch % 10 == 0: torch.save( model.state_dict(), os.path.join( cp_dir, 'checkpoint_{}_final.pth'.format(args.start_epoch + epoch)), ) torch.save( model.state_dict(), os.path.join( cp_dir, 'checkpoint_{}_final.pth'.format(args.start_epoch + args.max_epoch)), ) return None
def train(self): torch.cuda.empty_cache() ###################### # Save / Load model ###################### if self.opt.continue_train: try: self.continue_from_latest_checkpoint() except CyganException as e: self.logger.error(e) self.opt.continue_train = False self.reset_save() else: self.reset_save() self.add_file_logger() ###################### # Dataset ###################### if self.opt.model == 'base': dataset = SteelyDataset(self.opt.genreA, self.opt.genreB, self.opt.phase, use_mix=False) else: dataset = SteelyDataset(self.opt.genreA, self.opt.genreB, self.opt.phase, use_mix=True) dataset_size = len(dataset) iter_num = int(dataset_size / self.opt.batch_size) self.logger.info( f'Dataset loaded, genreA: {self.opt.genreA}, genreB: {self.opt.genreB}, total size: {dataset_size}.' ) ###################### # Initiate ###################### lambda_A = 10.0 # weight for cycle loss (A -> B -> A^) lambda_B = 10.0 # weight for cycle loss (B -> A -> B^) lambda_identity = 0.5 criterionGAN = GANLoss(gan_mode='lsgan') criterionCycle = nn.L1Loss() criterionIdt = nn.L1Loss() GLoss_meter = MovingAverageValueMeter(self.opt.plot_every) DLoss_meter = MovingAverageValueMeter(self.opt.plot_every) CycleLoss_meter = MovingAverageValueMeter(self.opt.plot_every) # loss meters losses = {} scores = {} losses_dict = {'loss_G': [], 'loss_D': [], 'loss_C': [], 'epoch': []} ###################### # Start Training ###################### for epoch in range(self.opt.start_epoch, self.opt.max_epoch): loader = DataLoader(dataset, batch_size=self.opt.batch_size, shuffle=True, num_workers=self.opt.num_threads, drop_last=True) epoch_start_time = time.time() for i, data in enumerate(loader): real_A = torch.unsqueeze(data[:, 0, :, :], 1).to(self.device, dtype=torch.float) real_B = torch.unsqueeze(data[:, 1, :, :], 1).to(self.device, dtype=torch.float) gaussian_noise = torch.abs( torch.normal(mean=torch.zeros(self.opt.data_shape), std=self.opt.gaussian_std)).to( self.device, dtype=torch.float) if self.opt.model == 'base': ###################### # Generator ###################### fake_B = self.generator_A2B(real_A) # X -> Y' fake_A = self.generator_B2A(real_B) # Y -> X' fake_B_copy = copy.copy(fake_B) fake_A_copy = copy.copy(fake_A) DB_fake = self.discriminator_B( fake_B + gaussian_noise) # netD_x provide feedback to netG_x DA_fake = self.discriminator_A(fake_A + gaussian_noise) loss_G_A2B = criterionGAN(DB_fake, True) loss_G_B2A = criterionGAN(DA_fake, True) # cycle_consistence cycle_A = self.generator_B2A(fake_B) # Y' -> X^ cycle_B = self.generator_A2B(fake_A) # Y -> X' -> Y^ loss_cycle_A2B = criterionCycle(cycle_A, real_A) * lambda_A loss_cycle_B2A = criterionCycle(cycle_B, real_B) * lambda_B # identity loss if lambda_identity > 0: # netG_x should be identity if real_y is fed: ||netG_x(real_y) - real_y|| idt_A = self.generator_A2B(real_B) idt_B = self.generator_B2A(real_A) loss_idt_A = criterionIdt( idt_A, real_B) * lambda_A * lambda_identity loss_idt_B = criterionIdt( idt_B, real_A) * lambda_A * lambda_identity else: loss_idt_A = 0. loss_idt_B = 0. loss_idt = loss_idt_A + loss_idt_B self.GA2B_optimizer.zero_grad( ) # set g_x and g_y gradients to zero loss_A2B = loss_G_A2B + loss_cycle_A2B + loss_idt_A loss_A2B.backward(retain_graph=True) self.GA2B_optimizer.step() self.GB2A_optimizer.zero_grad( ) # set g_x and g_y gradients to zero loss_B2A = loss_G_B2A + loss_cycle_B2A + loss_idt_B loss_B2A.backward(retain_graph=True) self.GB2A_optimizer.step() cycle_loss = loss_cycle_A2B + loss_cycle_B2A CycleLoss_meter.add(cycle_loss.item()) loss_G = loss_G_A2B + loss_G_B2A + loss_idt GLoss_meter.add(loss_G.item()) ###################### # Sample ###################### fake_A_sample, fake_B_sample = (None, None) if self.opt.use_image_pool: [fake_A_sample, fake_B_sample] = self.pool([fake_A_copy, fake_B_copy]) ###################### # Discriminator ###################### # loss_real DA_real = self.discriminator_A(real_A + gaussian_noise) DB_real = self.discriminator_B(real_B + gaussian_noise) loss_DA_real = criterionGAN(DA_real, True) loss_DB_real = criterionGAN(DB_real, True) # loss fake if self.opt.use_image_pool: DA_fake_sample = self.discriminator_A(fake_A_sample + gaussian_noise) DB_fake_sample = self.discriminator_B(fake_B_sample + gaussian_noise) loss_DA_fake = criterionGAN(DA_fake_sample, False) loss_DB_fake = criterionGAN(DB_fake_sample, False) else: loss_DA_fake = criterionGAN(DA_fake, False) loss_DB_fake = criterionGAN(DB_fake, False) # loss and backward self.DA_optimizer.zero_grad() loss_DA = (loss_DA_real + loss_DA_fake) * 0.5 loss_DA.backward() self.DA_optimizer.step() self.DB_optimizer.zero_grad() loss_DB = (loss_DB_real + loss_DB_fake) * 0.5 loss_DB.backward() self.DB_optimizer.step() loss_D = loss_DA + loss_DB DLoss_meter.add(loss_D.item()) else: real_mixed = torch.unsqueeze(data[:, 2, :, :], 1).to(self.device, dtype=torch.float) ###################### # Generator ###################### fake_B = self.generator_A2B(real_A) # X -> Y' fake_A = self.generator_B2A(real_B) # Y -> X' fake_B_copy = fake_B.detach().clone() fake_A_copy = fake_A.detach().clone() DB_fake = self.discriminator_B( fake_B + gaussian_noise) # netD_x provide feedback to netG_x DA_fake = self.discriminator_A(fake_A + gaussian_noise) loss_G_A2B = criterionGAN(DB_fake, True) loss_G_B2A = criterionGAN(DA_fake, True) # cycle_consistence cycle_A = self.generator_B2A(fake_B) # Y' -> X^ cycle_B = self.generator_A2B(fake_A) # Y -> X' -> Y^ loss_cycle_A2B = criterionCycle(cycle_A, real_A) * lambda_A loss_cycle_B2A = criterionCycle(cycle_B, real_B) * lambda_B # identity loss if lambda_identity > 0: # netG_x should be identity if real_y is fed: ||netG_x(real_y) - real_y|| idt_A = self.generator_A2B(real_B) idt_B = self.generator_B2A(real_A) loss_idt_A = criterionIdt( idt_A, real_B) * lambda_A * lambda_identity loss_idt_B = criterionIdt( idt_B, real_A) * lambda_A * lambda_identity else: loss_idt_A = 0. loss_idt_B = 0. loss_idt = loss_idt_A + loss_idt_B self.GA2B_optimizer.zero_grad( ) # set g_x and g_y gradients to zero loss_A2B = loss_G_A2B + loss_cycle_A2B + loss_idt_A loss_A2B.backward(retain_graph=True) self.GA2B_optimizer.step() self.GB2A_optimizer.zero_grad( ) # set g_x and g_y gradients to zero loss_B2A = loss_G_B2A + loss_cycle_B2A + loss_idt_B loss_B2A.backward(retain_graph=True) self.GB2A_optimizer.step() cycle_loss = loss_cycle_A2B + loss_cycle_B2A CycleLoss_meter.add(cycle_loss.item()) loss_G = loss_G_A2B + loss_G_B2A + loss_idt GLoss_meter.add(loss_G.item()) ###################### # Sample ###################### fake_A_sample, fake_B_sample = (None, None) if self.opt.use_image_pool: [fake_A_sample, fake_B_sample] = self.pool([fake_A_copy, fake_B_copy]) ###################### # Discriminator ###################### # loss_real DA_real = self.discriminator_A(real_A + gaussian_noise) DB_real = self.discriminator_B(real_B + gaussian_noise) DA_real_all = self.discriminator_A_all(real_mixed + gaussian_noise) DB_real_all = self.discriminator_B_all(real_mixed + gaussian_noise) loss_DA_real = criterionGAN(DA_real, True) loss_DB_real = criterionGAN(DB_real, True) loss_DA_all_real = criterionGAN(DA_real_all, True) loss_DB_all_real = criterionGAN(DB_real_all, True) # loss fake if self.opt.use_image_pool: DA_fake_sample = self.discriminator_A(fake_A_sample + gaussian_noise) DB_fake_sample = self.discriminator_B(fake_B_sample + gaussian_noise) DA_fake_sample_all = self.discriminator_A_all( fake_A_sample + gaussian_noise) DB_fake_sample_all = self.discriminator_B_all( fake_B_sample + gaussian_noise) loss_DA_all_fake = criterionGAN( DA_fake_sample_all, False) loss_DB_all_fake = criterionGAN( DB_fake_sample_all, False) loss_DA_fake = criterionGAN(DA_fake_sample, False) loss_DB_fake = criterionGAN(DB_fake_sample, False) else: DA_fake_all = self.discriminator_A_all(fake_A_copy + gaussian_noise) DB_fake_all = self.discriminator_B_all(fake_B_copy + gaussian_noise) loss_DA_all_fake = criterionGAN(DA_fake_all, False) loss_DB_all_fake = criterionGAN(DB_fake_all, False) loss_DA_fake = criterionGAN(DA_fake, False) loss_DB_fake = criterionGAN(DB_fake, False) # loss and backward self.DA_optimizer.zero_grad() loss_DA = (loss_DA_real + loss_DA_fake) * 0.5 loss_DA.backward() self.DA_optimizer.step() self.DB_optimizer.zero_grad() loss_DB = (loss_DB_real + loss_DB_fake) * 0.5 loss_DB.backward() self.DB_optimizer.step() self.DA_all_optimizer.zero_grad() loss_DA_all = (loss_DA_all_real + loss_DA_all_fake) * 0.5 loss_DA_all.backward() self.DA_all_optimizer.step() self.DB_all_optimizer.zero_grad() loss_DB_all = (loss_DB_all_real + loss_DB_all_fake) * 0.5 loss_DB_all.backward() self.DB_all_optimizer.step() loss_D = loss_DA + loss_DB + loss_DB_all + loss_DA_all DLoss_meter.add(loss_D.item()) ###################### # Snapshot ###################### if i % self.opt.plot_every == 0: file_name = self.opt.name + '_snap_%03d_%05d.png' % ( epoch, i, ) # test_path = os.path.join(self.opt.checkpoint_path, file_name) # tv.utils.save_image(fake_B, test_path, normalize=True) # self.logger.info(f'Snapshot {file_name} saved.') losses['loss_C'] = float(CycleLoss_meter.value()[0]) losses['loss_G'] = float(GLoss_meter.value()[0]) losses['loss_D'] = float(DLoss_meter.value()[0]) self.logger.info(str(losses)) self.logger.info('Epoch {} progress: {:.2%}\n'.format( epoch, i / iter_num)) # save model if epoch % self.opt.save_every == 0 or epoch == self.opt.max_epoch - 1: self.save_model(epoch) ###################### # lr_scheduler ###################### self.GA2B_scheduler.step(epoch) self.GB2A_scheduler.step(epoch) self.DA_scheduler.step(epoch) self.DB_scheduler.step(epoch) if self.opt.model != 'base': self.DA_all_scheduler.step(epoch) self.DB_all_scheduler.step(epoch) epoch_time = int(time.time() - epoch_start_time) ###################### # Logging ###################### self.logger.info( f'Epoch {epoch} finished, cost time {epoch_time}\n') self.logger.info(str(losses) + '\n\n') ###################### # Loss_Dict ###################### losses_dict['loss_C'].append(losses['loss_C']) losses_dict['loss_G'].append(losses['loss_G']) losses_dict['loss_D'].append(losses['loss_D']) losses_dict['epoch'].append(epoch) with open(self.opt.loss_save_path, 'w') as f: json.dump(losses_dict, f)
def train(): """Create the model and start the training.""" # === 1.Configuration print(CONFIG_PATH) # === select which GPU you want to use # === here assume to use 8 GPUs, idx are 0,1,2,3,...,7 os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, CONFIG.EXP.GPU_IDX)) device = get_device(torch.cuda.is_available()) cudnn.benchmark = True comment_init = "" writer = SummaryWriter(comment=comment_init) # Setup loss logger # === MovingAverageValueMeter(self,windowsize) # === - add(value): 记录value # === - reset() # === - value() : 返回MA和标准差 average_loss = MovingAverageValueMeter(CONFIG.SOLVER.AVERAGE_LOSS) if not os.path.exists(CONFIG.MODEL.SAVE_PATH): os.makedirs(CONFIG.MODEL.SAVE_PATH) # Path to save models checkpoint_dir = os.path.join( CONFIG.EXP.OUTPUT_DIR, # ./data "models", CONFIG.MODEL.NAME.lower(), # DeepLabV2_ResNet101_MSC CONFIG.DATASET.SPLIT.TRAIN, # train_aug ) # === checkpoint_dir: ./data/DeepLabV2_ResNet101_MSC/train_aug if not os.path.exists(checkpoint_dir): os.makedirs(checkpoint_dir) print("Checkpoint dst:", checkpoint_dir) # === 2.Dataloader === trainloader = data.DataLoader( VOCDataSet( CONFIG.DATASET.DIRECTORY, CONFIG.DATASET.LIST_PATH, max_iters=CONFIG.SOLVER.ITER_MAX * CONFIG.SOLVER.BATCH_SIZE.TRAIN, crop_size=(CONFIG.IMAGE.SIZE.TRAIN, CONFIG.IMAGE.SIZE.TRAIN), scale=CONFIG.DATASET.RANDOM.SCALE, mirror=CONFIG.DATASET.RANDOM.MIRROR, mean=IMG_MEAN, label_path=CONFIG.DATASET.SEG_LABEL), # for training batch_size=CONFIG.SOLVER.BATCH_SIZE.TRAIN, shuffle=True, num_workers=CONFIG.DATALOADER.NUM_WORKERS, pin_memory=True) # 使用iter(dataloader)返回的是一个迭代器,可以使用next访问 # loader_iter = iter(trainloader) # === 3.Create network & weights === print("Model:", CONFIG.MODEL.NAME) # model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.DATASET.N_CLASSES) model = DeepLabV2_DRN105_MSC(n_classes=CONFIG.DATASET.N_CLASSES) state_dict = torch.load(CONFIG.MODEL.INIT_MODEL) # model.base.load_state_dict(state_dict, strict=False) # to skip ASPP print(" Init:", CONFIG.MODEL.INIT_MODEL) # === show the skip weight for m in model.base.state_dict().keys(): if m not in state_dict.keys(): print(" Skip init:", m) # === DeepLabv2 = Res101+ASPP # === model.base = DeepLabv2 # === model = MSC(DeepLabv2) # model.base.load_state_dict(state_dict, # strict=False) # strict=False to skip ASPP model = nn.DataParallel(model) # multi-GPU model.to(device) # put in GPU is available # === 4.Loss definition criterion = nn.CrossEntropyLoss(ignore_index=CONFIG.DATASET.IGNORE_LABEL) criterion.to(device) # put in GPU is available # === 5.optimizer === optimizer = torch.optim.SGD( # cf lr_mult and decay_mult in train.prototxt params=[ { "params": get_params(model.module, key="1x"), "lr": CONFIG.SOLVER.LR, "weight_decay": CONFIG.SOLVER.WEIGHT_DECAY, }, { "params": get_params(model.module, key="10x"), "lr": 10 * CONFIG.SOLVER.LR, "weight_decay": CONFIG.SOLVER.WEIGHT_DECAY, }, { "params": get_params(model.module, key="20x"), "lr": 20 * CONFIG.SOLVER.LR, "weight_decay": 0.0, }, ], momentum=CONFIG.SOLVER.MOMENTUM, ) # Learning rate scheduler scheduler = PolynomialLR( optimizer=optimizer, step_size=CONFIG.SOLVER.LR_DECAY, iter_max=CONFIG.SOLVER.ITER_MAX, power=CONFIG.SOLVER.POLY_POWER, ) time_start = time.time() # set start time # === training iteration === for i_iter, batch in enumerate(trainloader, start=1): torch.set_grad_enabled(True) model.train() model.module.base.freeze_bn() optimizer.zero_grad() images, labels, _, _ = batch logits = model(images.to(device)) # <<<<<<<<<<<<<<<<<<<< # === Loss # === logits = [logits] + logits_pyramid + [logits_max] iter_loss = 0 loss = 0 for logit in logits: # Resize labels for {100%, 75%, 50%, Max} logits _, _, H, W = logit.shape labels_ = resize_labels(labels, size=(H, W)) iter_loss += criterion(logit, labels_.to(device)) # iter_loss /= CONFIG.SOLVER.ITER_SIZE iter_loss /= 4 iter_loss.backward() loss += float(iter_loss) average_loss.add(loss) # Update weights with accumulated gradients optimizer.step() # Update learning rate scheduler.step(epoch=i_iter) # TensorBoard writer.add_scalar("loss", average_loss.value()[0], global_step=i_iter) print( 'iter/max_iter = [{}/{}] completed, loss = {:4.3} time:{}'.format( i_iter, CONFIG.SOLVER.ITER_MAX, average_loss.value()[0], show_timing(time_start, time.time()))) # print('iter = ', i_iter, 'of', args.num_steps, '', # loss.data.cpu().numpy()) # === save final model if i_iter >= CONFIG.SOLVER.ITER_MAX: print('save final model as...{}'.format( osp.join(CONFIG.MODEL.SAVE_PATH, 'VOC12_' + str(CONFIG.SOLVER.ITER_MAX) + '.pth'))) torch.save( model.module.state_dict(), osp.join(CONFIG.MODEL.SAVE_PATH, 'VOC12_' + str(CONFIG.SOLVER.ITER_MAX) + '.pth')) break if i_iter % CONFIG.EXP.EVALUATE_ITER == 0: print("Evaluation....") evaluate_gpu(model, writer, i_iter) # === Save model every 250 iteration========================== # because DataParalel will add 'module' in each name of layer. # so here use model.module.state_dict() # ============================================================ if i_iter % CONFIG.MODEL.SAVE_EVERY_ITER == 0: print('saving model ...') torch.save( model.module.state_dict(), osp.join(CONFIG.MODEL.SAVE_PATH, 'VOC12_{}.pth'.format(i_iter)))
def main(config, cuda): device = torch.device("cuda" if cuda and torch.cuda.is_available() else "cpu") if cuda: current_device = torch.cuda.current_device() print("Running on", torch.cuda.get_device_name(current_device)) else: print("Running on CPU") # Configuration CONFIG = Dict(yaml.load(open(config))) # Dataset dataset = CocoStuff10k( root=CONFIG.ROOT, split="train", image_size=513, crop_size=CONFIG.IMAGE.SIZE.TRAIN, scale=True, flip=True, ) # DataLoader loader = torch.utils.data.DataLoader( dataset=dataset, batch_size=CONFIG.BATCH_SIZE, num_workers=CONFIG.NUM_WORKERS, shuffle=True, ) loader_iter = iter(loader) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES) state_dict = torch.load(CONFIG.INIT_MODEL) model.load_state_dict(state_dict, strict=False) # Skip "aspp" layer model = nn.DataParallel(model) model.to(device) # Optimizer optimizer = { "sgd": torch.optim.SGD( # cf lr_mult and decay_mult in train.prototxt params=[ { "params": get_lr_params(model.module, key="1x"), "lr": CONFIG.LR, "weight_decay": CONFIG.WEIGHT_DECAY, }, { "params": get_lr_params(model.module, key="10x"), "lr": 10 * CONFIG.LR, "weight_decay": CONFIG.WEIGHT_DECAY, }, { "params": get_lr_params(model.module, key="20x"), "lr": 20 * CONFIG.LR, "weight_decay": 0.0, }, ], momentum=CONFIG.MOMENTUM, ) }.get(CONFIG.OPTIMIZER) # Loss definition criterion = CrossEntropyLoss2d(ignore_index=CONFIG.IGNORE_LABEL) criterion.to(device) # TensorBoard Logger writer = SummaryWriter(CONFIG.LOG_DIR) loss_meter = MovingAverageValueMeter(20) model.train() model.module.scale.freeze_bn() for iteration in tqdm( range(1, CONFIG.ITER_MAX + 1), total=CONFIG.ITER_MAX, leave=False, dynamic_ncols=True, ): # Set a learning rate poly_lr_scheduler( optimizer=optimizer, init_lr=CONFIG.LR, iter=iteration - 1, lr_decay_iter=CONFIG.LR_DECAY, max_iter=CONFIG.ITER_MAX, power=CONFIG.POLY_POWER, ) # Clear gradients (ready to accumulate) optimizer.zero_grad() iter_loss = 0 for i in range(1, CONFIG.ITER_SIZE + 1): try: data, target = next(loader_iter) except: loader_iter = iter(loader) data, target = next(loader_iter) # Image data = data.to(device) # Propagate forward outputs = model(data) # Loss loss = 0 for output in outputs: # Resize target for {100%, 75%, 50%, Max} outputs target_ = resize_target(target, output.size(2)) target_ = target_.to(device) # Compute crossentropy loss loss += criterion(output, target_) # Backpropagate (just compute gradients wrt the loss) loss /= float(CONFIG.ITER_SIZE) loss.backward() iter_loss += float(loss) loss_meter.add(iter_loss) # Update weights with accumulated gradients optimizer.step() # TensorBoard if iteration % CONFIG.ITER_TF == 0: writer.add_scalar("train_loss", loss_meter.value()[0], iteration) for i, o in enumerate(optimizer.param_groups): writer.add_scalar("train_lr_group{}".format(i), o["lr"], iteration) # for name, param in model.named_parameters(): # name = name.replace('.', '/') # writer.add_histogram(name, param, iteration, bins="auto") # if param.requires_grad: # writer.add_histogram(name + '/grad', param.grad, iteration, bins="auto") # Save a model if iteration % CONFIG.ITER_SNAP == 0: torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, "checkpoint_{}.pth".format(iteration)), ) # Save a model if iteration % 100 == 0: torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, "checkpoint_current.pth"), ) torch.save( model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, "checkpoint_final.pth") )
def main(config, cuda): # Configuration with open(config) as f: CONFIG = yaml.load(f) cuda = cuda and torch.cuda.is_available() # Dataset dataset = get_dataset(CONFIG['DATASET'])( root=CONFIG['ROOT'], split='train', image_size=(CONFIG['IMAGE']['SIZE']['TRAIN'], CONFIG['IMAGE']['SIZE']['TRAIN']), scale=True, flip=True, # preload=True ) # DataLoader loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=CONFIG['BATCH_SIZE'], num_workers=CONFIG['NUM_WORKERS'], shuffle=True) loader_iter = iter(loader) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG['N_CLASSES']) state_dict = torch.load(CONFIG['INIT_MODEL']) model.load_state_dict(state_dict, strict=False) # Skip "aspp" layer if cuda: model.cuda() # Optimizer optimizer = { 'sgd': torch.optim.SGD( params=[ { 'params': get_1x_lr_params(model), 'lr': float(CONFIG['LR']) }, { 'params': get_10x_lr_params(model), 'lr': 10 * float(CONFIG['LR']) } # NOQA ], lr=float(CONFIG['LR']), momentum=float(CONFIG['MOMENTUM']), weight_decay=float(CONFIG['WEIGHT_DECAY'])), }.get(CONFIG['OPTIMIZER']) # Loss definition criterion = CrossEntropyLoss2d(ignore_index=CONFIG['IGNORE_LABEL']) if cuda: criterion.cuda() # TensorBoard Logger writer = SummaryWriter(CONFIG['LOG_DIR']) loss_meter = MovingAverageValueMeter(20) model.train() for iteration in tqdm(range(1, CONFIG['ITER_MAX'] + 1), total=CONFIG['ITER_MAX'], leave=False, dynamic_ncols=True): # Polynomial lr decay poly_lr_scheduler(optimizer=optimizer, init_lr=float(CONFIG['LR']), iter=iteration - 1, lr_decay_iter=CONFIG['LR_DECAY'], max_iter=CONFIG['ITER_MAX'], power=CONFIG['POLY_POWER']) optimizer.zero_grad() iter_loss = 0 for i in range(1, CONFIG['ITER_SIZE'] + 1): data, target = next(loader_iter) # Image data = data.cuda() if cuda else data data = Variable(data) # Forward propagation outputs = model(data) # Label target = resize_target(target, outputs[0].size(2)) target = target.cuda() if cuda else target target = Variable(target) # Aggregate losses for [100%, 75%, 50%, Max] loss = 0 for output in outputs: loss += criterion(output, target) loss /= CONFIG['ITER_SIZE'] iter_loss += loss.data[0] loss.backward() # Reload dataloader if ((iteration - 1) * CONFIG['ITER_SIZE'] + i) % len(loader) == 0: loader_iter = iter(loader) loss_meter.add(iter_loss) # Back propagation optimizer.step() # TensorBoard if iteration % CONFIG['ITER_TF'] == 0: writer.add_scalar('train_loss', loss_meter.value()[0], iteration) # Save a model if iteration % CONFIG['ITER_SNAP'] == 0: torch.save( model.state_dict(), osp.join(CONFIG['SAVE_DIR'], 'checkpoint_{}.pth.tar'.format(iteration))) # NOQA writer.add_text('log', 'Saved a model', iteration) torch.save(model.state_dict(), osp.join(CONFIG['SAVE_DIR'], 'checkpoint_final.pth.tar'))
def generateTB(self, period): self.writer = SummaryWriter(self.savepath + '/runs') self.loss_meter = MovingAverageValueMeter(20) self.tb = period
], momentum=CONFIG.SOLVER.MOMENTUM, ) # Learning rate scheduler scheduler = PolynomialLR( optimizer=optimizer, step_size=CONFIG.SOLVER.LR_DECAY, iter_max=CONFIG.SOLVER.ITER_MAX, power=CONFIG.SOLVER.POLY_POWER, ) # Setup loss logger writer = SummaryWriter(os.path.join("experiment", CONFIG.EXP_ID, "summary")) average_loss = MovingAverageValueMeter(CONFIG.SOLVER.AVERAGE_LOSS) # Path to save models checkpoint_dir = os.path.join("experiment", CONFIG.EXP_ID, "checkpoints") makedirs(checkpoint_dir) print("Checkpoint dst:", checkpoint_dir) # Random Dropout model.train() for iteration in tqdm( range(1, CONFIG.SOLVER.ITER_MAX + 1), total=CONFIG.SOLVER.ITER_MAX, dynamic_ncols=True, ):