def __init__(self): self.best_accuracy = 0.0 self.adjust_learning_rate = Config.adjust_learning_rate # all data self.data_train = MiniImageNetDataset.get_data_all(Config.data_root) self.task_train = MiniImageNetDataset(self.data_train, Config.num_way, Config.num_shot) self.task_train_loader = DataLoader(self.task_train, Config.batch_size, shuffle=True, num_workers=Config.num_workers) # model self.feature_encoder = RunnerTool.to_cuda(Config.feature_encoder) self.relation_network = RunnerTool.to_cuda(Config.relation_network) RunnerTool.to_cuda(self.feature_encoder.apply(RunnerTool.weights_init)) RunnerTool.to_cuda(self.relation_network.apply( RunnerTool.weights_init)) # optim self.feature_encoder_optim = torch.optim.SGD( self.feature_encoder.parameters(), lr=Config.learning_rate, momentum=0.9, weight_decay=5e-4) self.relation_network_optim = torch.optim.SGD( self.relation_network.parameters(), lr=Config.learning_rate, momentum=0.9, weight_decay=5e-4) # loss self.loss = RunnerTool.to_cuda(nn.MSELoss()) self.test_tool = TestTool(self.compare_fsl_test, data_root=Config.data_root, num_way=Config.num_way, num_shot=Config.num_shot, episode_size=Config.episode_size, test_episode=Config.test_episode, transform=self.task_train.transform_test) pass
def __init__(self): self.best_accuracy = 0.0 # all data self.data_train = MiniImageNetDataset.get_data_all(Config.data_root) self.task_train = MiniImageNetDataset(self.data_train, Config.num_way, Config.num_shot) self.task_train_loader = DataLoader(self.task_train, Config.batch_size, True, num_workers=Config.num_workers) # model self.feature_encoder = RunnerTool.to_cuda(Config.feature_encoder) self.relation_network = RunnerTool.to_cuda(Config.relation_network) RunnerTool.to_cuda(self.feature_encoder.apply(RunnerTool.weights_init)) RunnerTool.to_cuda(self.relation_network.apply( RunnerTool.weights_init)) # optim self.feature_encoder_optim = torch.optim.Adam( self.feature_encoder.parameters(), lr=Config.learning_rate) self.feature_encoder_scheduler = StepLR(self.feature_encoder_optim, Config.train_epoch // 3, gamma=0.5) self.relation_network_optim = torch.optim.Adam( self.relation_network.parameters(), lr=Config.learning_rate) self.relation_network_scheduler = StepLR(self.relation_network_optim, Config.train_epoch // 3, gamma=0.5) # loss self.loss = RunnerTool.to_cuda(nn.MSELoss()) self.test_tool = TestTool(self.compare_fsl_test, data_root=Config.data_root, num_way=Config.num_way, num_shot=Config.num_shot, episode_size=Config.episode_size, test_episode=Config.test_episode, transform=self.task_train.transform_test) pass
def train(self): Tools.print() Tools.print("Training...") for epoch in range(1, 1 + Config.train_epoch): self.feature_encoder.train() self.relation_network.train() Tools.print() fe_lr = self.adjust_learning_rate(self.feature_encoder_optim, epoch, Config.first_epoch, Config.t_epoch, Config.learning_rate) rn_lr = self.adjust_learning_rate(self.relation_network_optim, epoch, Config.first_epoch, Config.t_epoch, Config.learning_rate) Tools.print('Epoch: [{}] fe_lr={} rn_lr={}'.format( epoch, fe_lr, rn_lr)) all_loss = 0.0 for task_data, task_labels, task_index in tqdm( self.task_train_loader): task_data, task_labels = RunnerTool.to_cuda( task_data), RunnerTool.to_cuda(task_labels) # 1 calculate features relations = self.compare_fsl(task_data) # 2 loss loss = self.loss(relations, task_labels) all_loss += loss.item() # 3 backward self.feature_encoder.zero_grad() self.relation_network.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_( self.feature_encoder.parameters(), 0.5) torch.nn.utils.clip_grad_norm_( self.relation_network.parameters(), 0.5) self.feature_encoder_optim.step() self.relation_network_optim.step() ########################################################################### pass ########################################################################### # print Tools.print("{:6} loss:{:.3f}".format( epoch, all_loss / len(self.task_train_loader))) ########################################################################### ########################################################################### # Val if epoch % Config.val_freq == 0: self.feature_encoder.eval() self.relation_network.eval() Tools.print() Tools.print("Test {} .......".format(epoch)) val_accuracy = self.test_tool.val(episode=epoch, is_print=True) if val_accuracy > self.best_accuracy: self.best_accuracy = val_accuracy torch.save(self.feature_encoder.state_dict(), Config.fe_dir) torch.save(self.relation_network.state_dict(), Config.rn_dir) Tools.print("Save networks for epoch: {}".format(epoch)) pass pass ########################################################################### pass pass
def train(self): Tools.print() Tools.print("Training...") # Init Update try: self.feature_encoder.eval() self.relation_network.eval() self.ic_model.eval() Tools.print("Init label {} .......") self.produce_class.reset() for task_data, task_labels, task_index in tqdm( self.task_train_loader): ic_labels = RunnerTool.to_cuda(task_index[:, -1]) task_data, task_labels = RunnerTool.to_cuda( task_data), RunnerTool.to_cuda(task_labels) relations, query_features = self.compare_fsl(task_data) ic_out_logits, ic_out_l2norm = self.ic_model(query_features) self.produce_class.cal_label(ic_out_l2norm, ic_labels) pass Tools.print("Epoch: {}/{}".format(self.produce_class.count, self.produce_class.count_2)) finally: pass for epoch in range(Config.train_epoch): self.feature_encoder.train() self.relation_network.train() self.ic_model.train() Tools.print() fe_lr = self.adjust_learning_rate(self.feature_encoder_optim, epoch, Config.first_epoch, Config.t_epoch, Config.learning_rate) rn_lr = self.adjust_learning_rate(self.relation_network_optim, epoch, Config.first_epoch, Config.t_epoch, Config.learning_rate) ic_lr = self.adjust_learning_rate(self.ic_model_optim, epoch, Config.first_epoch, Config.t_epoch, Config.learning_rate) Tools.print('Epoch: [{}] fe_lr={} rn_lr={} ic_lr={}'.format( epoch, fe_lr, rn_lr, ic_lr)) self.produce_class.reset() all_loss, all_loss_fsl, all_loss_ic = 0.0, 0.0, 0.0 for task_data, task_labels, task_index in tqdm( self.task_train_loader): ic_labels = RunnerTool.to_cuda(task_index[:, -1]) task_data, task_labels = RunnerTool.to_cuda( task_data), RunnerTool.to_cuda(task_labels) ########################################################################### # 1 calculate features relations, query_features = self.compare_fsl(task_data) ic_out_logits, ic_out_l2norm = self.ic_model(query_features) # 2 ic_targets = self.produce_class.get_label(ic_labels) self.produce_class.cal_label(ic_out_l2norm, ic_labels) # 3 loss loss_fsl = self.fsl_loss(relations, task_labels) * Config.loss_fsl_ratio loss_ic = self.ic_loss(ic_out_logits, ic_targets) * Config.loss_ic_ratio loss = loss_fsl + loss_ic all_loss += loss.item() all_loss_fsl += loss_fsl.item() all_loss_ic += loss_ic.item() # 4 backward self.feature_encoder.zero_grad() self.relation_network.zero_grad() self.ic_model.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_( self.feature_encoder.parameters(), 0.5) torch.nn.utils.clip_grad_norm_( self.relation_network.parameters(), 0.5) torch.nn.utils.clip_grad_norm_(self.ic_model.parameters(), 0.5) self.feature_encoder_optim.step() self.relation_network_optim.step() self.ic_model_optim.step() ########################################################################### pass ########################################################################### # print Tools.print("{:6} loss:{:.3f} fsl:{:.3f} ic:{:.3f}".format( epoch + 1, all_loss / len(self.task_train_loader), all_loss_fsl / len(self.task_train_loader), all_loss_ic / len(self.task_train_loader))) Tools.print("Train: [{}] {}/{}".format(epoch, self.produce_class.count, self.produce_class.count_2)) ########################################################################### ########################################################################### # Val if epoch % Config.val_freq == 0: self.feature_encoder.eval() self.relation_network.eval() self.ic_model.eval() self.test_tool_ic.val(epoch=epoch) val_accuracy = self.test_tool_fsl.val(episode=epoch, is_print=True) if val_accuracy > self.best_accuracy: self.best_accuracy = val_accuracy torch.save(self.feature_encoder.state_dict(), Config.fe_dir) torch.save(self.relation_network.state_dict(), Config.rn_dir) torch.save(self.ic_model.state_dict(), Config.ic_dir) Tools.print("Save networks for epoch: {}".format(epoch)) pass pass ########################################################################### pass pass
def __init__(self): self.best_accuracy = 0.0 self.adjust_learning_rate = Config.adjust_learning_rate # all data self.data_train = MiniImageNetDataset.get_data_all(Config.data_root) self.task_train = MiniImageNetDataset(self.data_train, Config.num_way, Config.num_shot) self.task_train_loader = DataLoader(self.task_train, Config.batch_size, True, num_workers=Config.num_workers) # IC self.produce_class = ProduceClass(len(self.data_train), Config.ic_out_dim, Config.ic_ratio) self.produce_class.init() self.task_train.set_samples_class(self.produce_class.classes) # model self.feature_encoder = RunnerTool.to_cuda(Config.feature_encoder) self.relation_network = RunnerTool.to_cuda(Config.relation_network) self.ic_model = RunnerTool.to_cuda( ICModel(in_dim=Config.ic_in_dim, out_dim=Config.ic_out_dim)) RunnerTool.to_cuda(self.feature_encoder.apply(RunnerTool.weights_init)) RunnerTool.to_cuda(self.relation_network.apply( RunnerTool.weights_init)) RunnerTool.to_cuda(self.ic_model.apply(RunnerTool.weights_init)) # optim self.feature_encoder_optim = torch.optim.SGD( self.feature_encoder.parameters(), lr=Config.learning_rate, momentum=0.9, weight_decay=5e-4) self.relation_network_optim = torch.optim.SGD( self.relation_network.parameters(), lr=Config.learning_rate, momentum=0.9, weight_decay=5e-4) self.ic_model_optim = torch.optim.SGD(self.ic_model.parameters(), lr=Config.learning_rate, momentum=0.9, weight_decay=5e-4) # loss self.ic_loss = RunnerTool.to_cuda(nn.CrossEntropyLoss()) self.fsl_loss = RunnerTool.to_cuda(nn.MSELoss()) # Eval self.test_tool_fsl = TestTool(self.compare_fsl_test, data_root=Config.data_root, num_way=Config.num_way, num_shot=Config.num_shot, episode_size=Config.episode_size, test_episode=Config.test_episode, transform=self.task_train.transform_test) self.test_tool_ic = ICTestTool(feature_encoder=self.feature_encoder, ic_model=self.ic_model, data_root=Config.data_root, batch_size=Config.batch_size, num_workers=Config.num_workers, ic_out_dim=Config.ic_out_dim) pass
def __init__(self): self.best_accuracy = 0.0 # all data self.data_train = MiniImageNetDataset.get_data_all(Config.data_root) self.task_train = MiniImageNetDataset(self.data_train, Config.num_way, Config.num_shot) self.task_train_loader = DataLoader(self.task_train, Config.batch_size, True, num_workers=Config.num_workers) # IC self.produce_class = ProduceClass(len(self.data_train), Config.ic_out_dim, Config.ic_ratio) self.produce_class.init() # model self.feature_encoder = RunnerTool.to_cuda(Config.feature_encoder) self.relation_network = RunnerTool.to_cuda(Config.relation_network) self.ic_model = RunnerTool.to_cuda(Config.ic_model) RunnerTool.to_cuda(self.feature_encoder.apply(RunnerTool.weights_init)) RunnerTool.to_cuda(self.relation_network.apply( RunnerTool.weights_init)) RunnerTool.to_cuda(self.ic_model.apply(RunnerTool.weights_init)) # optim self.feature_encoder_optim = torch.optim.Adam( self.feature_encoder.parameters(), lr=Config.learning_rate) self.relation_network_optim = torch.optim.Adam( self.relation_network.parameters(), lr=Config.learning_rate) self.ic_model_optim = torch.optim.Adam(self.ic_model.parameters(), lr=Config.learning_rate) self.feature_encoder_scheduler = StepLR(self.feature_encoder_optim, Config.train_epoch // 3, gamma=0.5) self.relation_network_scheduler = StepLR(self.relation_network_optim, Config.train_epoch // 3, gamma=0.5) self.ic_model_scheduler = StepLR(self.ic_model_optim, Config.train_epoch // 3, gamma=0.5) # loss self.ic_loss = RunnerTool.to_cuda(nn.CrossEntropyLoss()) self.fsl_loss = RunnerTool.to_cuda(nn.MSELoss()) # Eval self.test_tool_fsl = TestTool(self.compare_fsl_test, data_root=Config.data_root, num_way=Config.num_way, num_shot=Config.num_shot, episode_size=Config.episode_size, test_episode=Config.test_episode, transform=self.task_train.transform_test) self.test_tool_ic = ICTestTool(feature_encoder=self.feature_encoder, ic_model=self.ic_model, data_root=Config.data_root, batch_size=Config.batch_size, num_workers=Config.num_workers, ic_out_dim=Config.ic_out_dim) pass
def __init__(self): self.feature_encoder = RunnerTool.to_cuda(Config.feature_encoder) self.relation_network = RunnerTool.to_cuda(Config.relation_network) pass