def __init__(self, experiment, device): config_file = os.path.join(CONFIG_DIR, experiment + '.json') assert os.path.exists( config_file), 'config file {} does not exist'.format(config_file) self.experiment = experiment with open(config_file, 'r') as f: configs = json.load(f) self.device = int(device) self.lr = configs['lr'] self.max_epochs = configs['max-epochs'] self.train_batch_size = configs['train-batch-size'] self.test_batch_size = configs['test-batch-size'] self.n_epochs = 0 self.log_dir = os.path.join(LOG_DIR, experiment) if not os.path.exists(self.log_dir): os.makedirs(self.log_dir) self.tboard_writer = tensorboardX.SummaryWriter(log_dir=self.log_dir) self.checkpoint_dir = os.path.join(CHECKPOINT_DIR, experiment) if not os.path.exists(self.checkpoint_dir): os.makedirs(self.checkpoint_dir) self.model = mstcn.MultiStageModel( num_stages=N_STAGES, num_layers=N_LAYERS, num_f_maps=N_FEATURE_MAPS, dim=IN_CHANNELS, num_classes=breakfast.N_MSTCN_CLASSES) self.model = self.model.cuda(self.device) self.ce_loss = nn.CrossEntropyLoss(ignore_index=-100).cuda(self.device) self.mse_loss = nn.MSELoss(reduction='none').cuda(self.device) if configs['optim'] == 'adam': self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr) elif configs['optim'] == 'sgd': self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr, momentum=configs['momentum'], nesterov=configs['nesterov']) else: raise ValueError('no such optimizer') if configs['scheduler'] == 'step': self.scheduler = optim.lr_scheduler.StepLR( self.optimizer, step_size=configs['lr-step'], gamma=configs['lr-decay']) elif configs['scheduler'] == 'plateau': self.scheduler = optim.lr_scheduler.ReduceLROnPlateau( self.optimizer, mode='min', patience=configs['lr-step']) else: raise ValueError('no such scheduler') self._load_checkpoint() self.submission_dir = os.path.join(SUBMISSION_DIR, self.experiment)
def __init__(self, experiment, device): config_file = os.path.join(CONFIG_DIR, experiment + '.json') assert os.path.exists( config_file), 'config file {} does not exist'.format(config_file) self.experiment = experiment with open(config_file, 'r') as f: configs = json.load(f) self.device = int(device) self.gen_lr = configs['gen-lr'] self.dis_lr = configs['dis-lr'] self.max_epochs = configs['max-epochs'] self.train_batch_size = configs['train-batch-size'] self.test_batch_size = configs['test-batch-size'] self.n_epochs = 0 self.log_dir = os.path.join(LOG_DIR, experiment) if not os.path.exists(self.log_dir): os.makedirs(self.log_dir) self.tboard_writer = tensorboardX.SummaryWriter(log_dir=self.log_dir) self.checkpoint_dir = os.path.join(CHECKPOINT_DIR, experiment) if not os.path.exists(self.checkpoint_dir): os.makedirs(self.checkpoint_dir) self.source_to_target_gen = cyclegan.VideoFeatGenerator( n_layers=10, in_channels=IN_CHANNELS, reduction_factor=4).cuda(self.device) self.target_to_source_gen = cyclegan.VideoFeatGenerator( n_layers=10, in_channels=IN_CHANNELS, reduction_factor=4).cuda(self.device) self.source_dis = cyclegan.VideoFeatDiscriminator( num_layers=N_LAYERS, num_f_maps=N_FEATURE_MAPS, dim=IN_CHANNELS).cuda(self.device) self.target_dis = cyclegan.VideoFeatDiscriminator( num_layers=N_LAYERS, num_f_maps=N_FEATURE_MAPS, dim=IN_CHANNELS).cuda(self.device) self.mstcn_model = mstcn.MultiStageModel(num_stages=N_STAGES, num_layers=N_LAYERS, num_f_maps=N_FEATURE_MAPS, dim=IN_CHANNELS, num_classes=breakfast.N_MSTCN_CLASSES)\ .cuda(self.device) self.mstcn_model.eval() self.mstcn_model_config = configs['mstcn-config'] self._load_mstcn_model() self.ce_loss = nn.CrossEntropyLoss(ignore_index=-100).cuda(self.device) self.mse_loss = nn.MSELoss(reduction='none').cuda(self.device) self.reconstruction_loss_fn = nn.L1Loss() gen_parameters = list(self.target_to_source_gen.parameters()) + list( self.source_to_target_gen.parameters()) dis_parameters = list(self.target_dis.parameters()) + list( self.source_dis.parameters()) if configs['optim'] == 'adam': self.gen_optimizer = optim.Adam(gen_parameters, lr=self.gen_lr) self.dis_optimizer = optim.Adam(dis_parameters, lr=self.dis_lr) elif configs['optim'] == 'sgd': self.gen_optimizer = optim.SGD(gen_parameters, lr=self.gen_lr, momentum=configs['momentum'], nesterov=configs['nesterov']) self.dis_optimizer = optim.SGD(dis_parameters, lr=self.gen_lr, momentum=configs['momentum'], nesterov=configs['nesterov']) else: raise ValueError('no such optimizer') if configs['scheduler'] == 'step': self.gen_scheduler = optim.lr_scheduler.StepLR( self.gen_optimizer, step_size=configs['gen-lr-step'], gamma=configs['gen-lr-decay']) self.dis_scheduler = optim.lr_scheduler.StepLR( self.dis_optimizer, step_size=configs['dis-lr-step'], gamma=configs['dis-lr-decay']) elif configs['scheduler'] == 'plateau': self.gen_scheduler = optim.lr_scheduler.ReduceLROnPlateau( self.gen_optimizer, mode='min', patience=configs['gen-lr-step']) self.dis_scheduler = optim.lr_scheduler.ReduceLROnPlateau( self.dis_optimizer, mode='min', patience=configs['dis-lr-step']) else: raise ValueError('no such scheduler') self._load_checkpoint()