def __init__(self, dataset_class, txt_file, root_dir, transform, num_workers): super(TrainClass, self).__init__() self.args = args self.device = torch.device("cuda:%d" % args.cuda if torch.cuda.is_available() else "cpu") self.writer = SummaryWriter('log_three') self.model = Stage2Model().to(self.device) if args.optimizer == 'Adam': self.optimizer = [optim.Adam(self.model.parameters(), self.args.lr) for _ in range(4)] else: self.optimizer = [optim.SGD(self.model.parameters(), self.args.lr, momentum=self.args.momentum, weight_decay=self.args.weight_decay) for _ in range(4)] self.criterion = [nn.CrossEntropyLoss() for _ in range(4)] # self.metric = nn.CrossEntropyLoss() self.metric = [F1Accuracy() for _ in range(4)] self.train_loader = None self.eval_loader = None self.ckpt_dir = "checkpoint_%s" % uuid self.display_freq = args.display_freq self.scheduler = [optim.lr_scheduler.StepLR(self.optimizer[i], step_size=5, gamma=0.5) for i in range(4)] self.best_error = [float('Inf'), float('Inf'), float('Inf'), float('Inf')] self.best_accu = [float('-Inf'), float('-Inf'), float('-Inf'), float('-Inf')] self.load_dataset(dataset_class, txt_file, root_dir, transform, num_workers)
def __init__(self, batch_size, is_shuffle, num_workers, state_files): super(Simple_two, self).__init__() self.args = None self.get_args() self.device = torch.device("cuda:%d" % self.args.cuda if torch.cuda.is_available() else "cpu") self.model1 = Stage1Model().to(self.device) self.model2 = Stage2Model().to(self.device) self.reverse = ReverseTModel().to(self.device) self.root_dir = "/data1/yinzi/datas" self.predict1 = None self.predict2 = None self.all_predict = None self.best_error = float('Inf') self.F1_name_list = ['eyebrow1', 'eyebrow2', 'eye1', 'eye2', 'nose', 'u_lip', 'i_mouth', 'l_lip'] self.TP = {x: 0.0 + 1e-20 for x in self.F1_name_list} self.FP = {x: 0.0 + 1e-20 for x in self.F1_name_list} self.TN = {x: 0.0 + 1e-20 for x in self.F1_name_list} self.FN = {x: 0.0 + 1e-20 for x in self.F1_name_list} self.recall = {x: 0.0 + 1e-20 for x in self.F1_name_list} self.precision = {x: 0.0 + 1e-20 for x in self.F1_name_list} self.F1_list = {x: [] for x in self.F1_name_list} self.F1 = {x: 0.0 + 1e-20 for x in self.F1_name_list} self.recall_overall_list = {x: [] for x in self.F1_name_list} self.precision_overall_list = {x: [] for x in self.F1_name_list} self.recall_overall = 0.0 self.precision_overall = 0.0 self.F1_overall = 0.0 self.dataset = None self.dataloader = None self.batch_size = batch_size self.is_shuffle = is_shuffle self.num_workers = num_workers self.state_files = state_files self.map_location = self.device self.stage2_loss_func = torch.nn.CrossEntropyLoss() self.optim2 = optim.Adam(self.model2.parameters(), self.args.lr) self.step = 0 self.epoch = 0 self.ckpt_dir = "checkpoints_%s" % uuid self.writer = SummaryWriter("logs") self.get_dataloader()
def __init__(self, argus=args): super(TrainModel, self).__init__() self.label_channels = 9 # ============== not neccessary =============== self.train_logger = None self.eval_logger = None self.args = argus # ============== neccessary =============== self.writer = SummaryWriter('log') self.step = 0 self.epoch = 0 self.best_error = float('Inf') self.best_accu = float('-Inf') self.device = torch.device( "cuda:%d" % args.cuda if torch.cuda.is_available() else "cpu") self.model = Stage2Model().to(self.device) self.model1 = Stage1Model().to(self.device) # self.reverse = ReverseTModel().to(self.device) # self.optimizer = optim.SGD(self.model.parameters(), self.args.lr, momentum=0.9, weight_decay=0.0) self.optimizer = optim.Adam(self.model.parameters(), self.args.lr) self.criterion = nn.CrossEntropyLoss() # self.criterion = nn.BCEWithLogitsLoss() # self.metric = nn.CrossEntropyLoss() self.metric = F1Accuracy() self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, step_size=5, gamma=0.5) self.train_loader = stage1_dataloaders['train'] self.eval_loader = stage1_dataloaders['val'] self.ckpt_dir = "checkpoints_%s" % uuid self.display_freq = args.display_freq # call it to check all members have been intiated self.check_init()
def __init__(self, dataset_class, txt_file, root_dir, transform, num_workers): super(TrainMGPU, self).__init__(dataset_class, txt_file, root_dir, transform, num_workers) self.model = nn.DataParallel(Stage2Model(), device_ids=[0, 1, 2, 3, 4]) self.model = self.model.to(self.device)