def main(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # * Step 1: init data folders print("init data folders") # * Init character folders for dataset construction metatrain_character_folders, metatest_character_folders = tg.mini_imagenet_folders( ) # * Step 2: init neural networks print("init neural networks") feature_encoder = models.CNNEncoder() actor = models.Actor(FEATURE_DIM, RELATION_DIM, CLASS_NUM) critic = models.Critic(FEATURE_DIM, RELATION_DIM) #feature_encoder = torch.nn.DataParallel(feature_encoder) #actor = torch.nn.DataParallel(actor) #critic = torch.nn.DataParallel(critic) feature_encoder.train() actor.train() critic.train() feature_encoder.apply(models.weights_init) actor.apply(models.weights_init) critic.apply(models.weights_init) feature_encoder.to(device) actor.to(device) critic.to(device) agent = a2cAgent.A2CAgent(actor, critic, GAMMA, ENTROPY_WEIGHT, FEATURE_DIM, RELATION_DIM, CLASS_NUM, device) #feature_encoder.eval() #relation_network.eval() if os.path.exists( str("./models/miniimagenet_feature_encoder_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl")): feature_encoder.load_state_dict( torch.load( str("./models/miniimagenet_feature_encoder_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl"))) print("load feature encoder success") if os.path.exists( str("./models/miniimagenet_actor_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl")): actor.load_state_dict( torch.load( str("./models/miniimagenet_actor_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl"))) print("load actor network success") if os.path.exists( str("./models/miniimagenet_critic_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl")): critic.load_state_dict( torch.load( str("./models/miniimagenet_critic_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl"))) print("load critic network success") max_accuracy_list = [] mean_accuracy_list = [] for episode in range(1): total_accuracy = [] for i in range(TEST_EPISODE): # * Generate env env_states_list = [] env_labels_list = [] number_of_query_image = 15 task = tg.MiniImagenetTask(metatest_character_folders, CLASS_NUM, SAMPLE_NUM_PER_CLASS, number_of_query_image) sample_dataloader = tg.get_mini_imagenet_data_loader( task, num_per_class=SAMPLE_NUM_PER_CLASS, split="train", shuffle=False) test_dataloader = tg.get_mini_imagenet_data_loader( task, num_per_class=number_of_query_image, split="test", shuffle=True) sample_images, sample_labels = next(iter(sample_dataloader)) test_images, test_labels = next(iter(test_dataloader)) sample_images, sample_labels = sample_images.to( device), sample_labels.to(device) test_images, test_labels = test_images.to(device), test_labels.to( device) # * calculate features sample_features = feature_encoder(sample_images) sample_features = sample_features.view(CLASS_NUM, SAMPLE_NUM_PER_CLASS, FEATURE_DIM, 19, 19) sample_features = torch.sum(sample_features, 1).squeeze(1) test_features = feature_encoder(test_images) # * calculate relations # * each batch sample link to every samples to calculate relations # * to form a 100x128 matrix for relation network sample_features_ext = sample_features.unsqueeze(0).repeat( number_of_query_image * CLASS_NUM, 1, 1, 1, 1) test_features_ext = test_features.unsqueeze(0).repeat( CLASS_NUM, 1, 1, 1, 1) test_features_ext = torch.transpose(test_features_ext, 0, 1) relation_pairs = torch.cat( (sample_features_ext, test_features_ext), 2).view(-1, FEATURE_DIM * 2, 19, 19) env_states_list.append(relation_pairs) env_labels_list.append(test_labels) test_env = a2cAgent.env(env_states_list, env_labels_list) rewards = agent.test(test_env) test_accuracy = rewards / len(test_labels) print(test_accuracy) total_accuracy.append(test_accuracy) mean_accuracy, conf_int = mean_confidence_interval(total_accuracy) print(f"Total accuracy : {mean_accuracy:.4f}") print(f"confidence interval : {conf_int:.4f}")
def main(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # * Step 1: init data folders print("init data folders") # * Init character folders for dataset construction metatrain_character_folders, metatest_character_folders = tg.omniglot_character_folders() # * Step 2: init neural networks print("init neural networks") feature_encoder = models.CNNEncoder() actor = models.Actor(FEATURE_DIM, RELATION_DIM, CLASS_NUM) critic = models.Critic(FEATURE_DIM, RELATION_DIM) #feature_encoder = torch.nn.DataParallel(feature_encoder) #actor = torch.nn.DataParallel(actor) #critic = torch.nn.DataParallel(critic) feature_encoder.train() actor.train() critic.train() feature_encoder.apply(models.weights_init) actor.apply(models.weights_init) critic.apply(models.weights_init) feature_encoder.to(device) actor.to(device) critic.to(device) cross_entropy = nn.CrossEntropyLoss() feature_encoder_optim = torch.optim.Adam(feature_encoder.parameters(), lr=LEARNING_RATE) feature_encoder_scheduler = StepLR(feature_encoder_optim, step_size=10000, gamma=0.5) actor_optim = torch.optim.Adam(actor.parameters(), lr=LEARNING_RATE) actor_scheduler = StepLR(actor_optim, step_size=10000, gamma=0.5) critic_optim = torch.optim.Adam(critic.parameters(), lr=LEARNING_RATE * 10) critic_scheduler = StepLR(critic_optim, step_size=10000, gamma=0.5) agent = a2cAgent.A2CAgent(actor, critic, GAMMA, ENTROPY_WEIGHT, FEATURE_DIM, RELATION_DIM, CLASS_NUM, device) if os.path.exists(str("./models/miniimagenet_feature_encoder_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl")): feature_encoder.load_state_dict(torch.load(str("./models/miniimagenet_feature_encoder_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl"))) print("load feature encoder success") if os.path.exists(str("./models/miniimagenet_actor_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl")): actor.load_state_dict(torch.load(str("./models/miniimagenet_actor_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl"))) print("load actor network success") if os.path.exists(str("./models/miniimagenet_critic_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl")): critic.load_state_dict(torch.load(str("./models/miniimagenet_critic_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl"))) print("load critic network success") # * Step 3: build graph print("Training...") last_accuracy = 0.0 mbal_loss_list = [] mbcl_loss_list = [] loss_list = [] number_of_query_image = 15 for episode in range(EPISODE): #print(f"EPISODE : {episode}") policy_losses = [] value_losses = [] for meta_batch in range(META_BATCH_RANGE): meta_env_states_list = [] meta_env_labels_list = [] for inner_batch in range(INNER_BATCH_RANGE): # * Generate environment env_states_list = [] env_labels_list = [] for env_num in range(ENV_LENGTH): task = tg.OmniglotTask(metatrain_character_folders, CLASS_NUM, SAMPLE_NUM_PER_CLASS, number_of_query_image) sample_dataloader = tg.get_data_loader(task, num_per_class=SAMPLE_NUM_PER_CLASS, split="train", shuffle=False) batch_dataloader = tg.get_data_loader(task, num_per_class=5, split="test", shuffle=True) samples, sample_labels = next(iter(sample_dataloader)) samples, sample_labels = samples.to(device), sample_labels.to(device) for batches, batch_labels in batch_dataloader: batches, batch_labels = batches.to(device), batch_labels.to(device) inner_sample_features = feature_encoder(samples) inner_sample_features = inner_sample_features.view(CLASS_NUM, SAMPLE_NUM_PER_CLASS, FEATURE_DIM, 5, 5) inner_sample_features = torch.sum(inner_sample_features, 1).squeeze(1) inner_batch_features = feature_encoder(batches) inner_sample_feature_ext = inner_sample_features.unsqueeze(0).repeat(5 * CLASS_NUM, 1, 1, 1, 1) inner_batch_features_ext = inner_batch_features.unsqueeze(0).repeat(CLASS_NUM, 1, 1, 1, 1) inner_batch_features_ext = torch.transpose(inner_batch_features_ext, 0, 1) inner_relation_pairs = torch.cat((inner_sample_feature_ext, inner_batch_features_ext), 2).view(-1, FEATURE_DIM * 2, 5, 5) env_states_list.append(inner_relation_pairs) env_labels_list.append(batch_labels) inner_env = a2cAgent.env(env_states_list, env_labels_list) agent.train(inner_env, inner_update=True) task = tg.OmniglotTask(metatrain_character_folders, CLASS_NUM, SAMPLE_NUM_PER_CLASS, number_of_query_image) sample_dataloader = tg.get_data_loader(task, num_per_class=SAMPLE_NUM_PER_CLASS, split="train", shuffle=False) batch_dataloader = tg.get_data_loader(task, num_per_class=number_of_query_image, split="test", shuffle=True) samples, sample_labels = next(iter(sample_dataloader)) samples, sample_labels = samples.to(device), sample_labels.to(device) # * Generate env for meta update # * init dataset # * sample_dataloader is to obtain previous samples for compare # * batch_dataloader is to batch samples for training # * num_per_class : number of query images # * sample datas batches, batch_labels = next(iter(batch_dataloader)) batches, batch_labels = batches.to(device), batch_labels.to(device) # * calculates features #feature_encoder.weight = feature_fast_weights sample_features = feature_encoder(samples) sample_features = sample_features.view(CLASS_NUM, SAMPLE_NUM_PER_CLASS, FEATURE_DIM, 5, 5) sample_features = torch.sum(sample_features, 1).squeeze(1) batch_features = feature_encoder(batches) # * calculate relations # * each batch sample link to every samples to calculate relations # * to form a 100 * 128 matrix for relation network sample_features_ext = sample_features.unsqueeze(0).repeat(number_of_query_image * CLASS_NUM, 1, 1, 1, 1) batch_features_ext = batch_features.unsqueeze(0).repeat(CLASS_NUM, 1, 1, 1, 1) batch_features_ext = torch.transpose(batch_features_ext, 0, 1) relation_pairs = torch.cat((sample_features_ext, batch_features_ext), 2).view(-1, FEATURE_DIM * 2, 5, 5) meta_env_states_list.append(relation_pairs) meta_env_labels_list.append(batch_labels) meta_env = a2cAgent.env(meta_env_states_list, meta_env_labels_list) agent.train(meta_env, policy_loss_list=policy_losses, value_loss_list=value_losses) feature_encoder_optim.zero_grad() actor_optim.zero_grad() critic_optim.zero_grad() torch.nn.utils.clip_grad_norm_(feature_encoder.parameters(), 0.5) torch.nn.utils.clip_grad_norm_(actor.parameters(), 0.5) torch.nn.utils.clip_grad_norm_(critic.parameters(), 0.5) meta_batch_actor_loss = torch.stack(policy_losses).mean() meta_batch_critic_loss = torch.stack(value_losses).mean() meta_batch_actor_loss.backward(retain_graph=True) meta_batch_critic_loss.backward() feature_encoder_optim.step() actor_optim.step() critic_optim.step() feature_encoder_scheduler.step(episode) actor_scheduler.step(episode) critic_scheduler.step(episode) if (episode + 1) % 100 == 0: mbal = meta_batch_actor_loss.cpu().detach().numpy() mbcl = meta_batch_critic_loss.cpu().detach().numpy() print(f"episode : {episode+1}, meta_batch_actor_loss : {mbal:.4f}, meta_batch_critic_loss : {mbcl:.4f}") mbal_loss_list.append(mbal) mbcl_loss_list.append(mbcl) loss_list.append(mbal + mbcl) if (episode + 1) % 500 == 0: print("Testing...") total_reward = 0 total_test_samples = 0 for i in range(TEST_EPISODE): # * Generate env env_states_list = [] env_labels_list = [] number_of_query_image = 10 task = tg.OmniglotTask(metatest_character_folders, CLASS_NUM, SAMPLE_NUM_PER_CLASS, number_of_query_image) sample_dataloader = tg.get_data_loader(task, num_per_class=SAMPLE_NUM_PER_CLASS, split="train", shuffle=False) test_dataloader = tg.get_data_loader(task, num_per_class=number_of_query_image, split="test", shuffle=True) sample_images, sample_labels = next(iter(sample_dataloader)) sample_images, sample_labels = sample_images.to(device), sample_labels.to(device) test_images, test_labels = next(iter(test_dataloader)) total_test_samples += len(test_labels) test_images, test_labels = test_images.to(device), test_labels.to(device) # * calculate features sample_features = feature_encoder(sample_images) sample_features = sample_features.view(CLASS_NUM, SAMPLE_NUM_PER_CLASS, FEATURE_DIM, 5, 5) sample_features = torch.sum(sample_features, 1).squeeze(1) test_features = feature_encoder(test_images) # * calculate relations # * each batch sample link to every samples to calculate relations # * to form a 100x128 matrix for relation network sample_features_ext = sample_features.unsqueeze(0).repeat(number_of_query_image * CLASS_NUM, 1, 1, 1, 1) test_features_ext = test_features.unsqueeze(0).repeat(CLASS_NUM, 1, 1, 1, 1) test_features_ext = torch.transpose(test_features_ext, 0, 1) relation_pairs = torch.cat((sample_features_ext, test_features_ext), 2).view(-1, FEATURE_DIM * 2, 5, 5) env_states_list.append(relation_pairs) env_labels_list.append(test_labels) test_env = a2cAgent.env(env_states_list, env_labels_list) rewards = agent.test(test_env) total_reward += rewards test_accuracy = total_reward / (1.0 * total_test_samples) mean_loss = np.mean(loss_list) mean_actor_loss = np.mean(mbal_loss_list) mean_critic_loss = np.mean(mbcl_loss_list) print(f'mean loss : {mean_loss}') print("test accuracy : ", test_accuracy) writer.add_scalar('1.loss', mean_loss, episode + 1) writer.add_scalar('2.mean_actor_loss', mean_actor_loss, episode + 1) writer.add_scalar('3.mean_critic_loss', mean_critic_loss, episode + 1) writer.add_scalar('4.test accuracy', test_accuracy, episode + 1) loss_list = [] mbal_loss_list = [] mbcl_loss_list = [] if test_accuracy > last_accuracy: # save networks torch.save( feature_encoder.state_dict(), str("./models/miniimagenet_feature_encoder_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl") ) torch.save( actor.state_dict(), str("./models/miniimagenet_actor_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl") ) torch.save( critic.state_dict(), str("./models/miniimagenet_critic_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl") ) print("save networks for episode:", episode) last_accuracy = test_accuracy
def main(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # * Step 1: init data folders print("init data folders") # * Init character folders for dataset construction metatrain_character_folders, metatest_character_folders = tg.omniglot_character_folders() # * Step 2: init neural networks print("init neural networks") feature_encoder = models.CNNEncoder() #actor = models.Actor(FEATURE_DIM, RELATION_DIM, CLASS_NUM) #critic = models.Critic(FEATURE_DIM, RELATION_DIM) agent = a2cAgent.A2CAgent(GAMMA, ENTROPY_WEIGHT, FEATURE_DIM, RELATION_DIM, CLASS_NUM, device) feature_encoder.apply(models.weights_init) #actor.apply(models.weights_init) #critic.apply(models.weights_init) feature_encoder.to(device) #actor.to(device) #critic.to(device) cross_entropy = nn.CrossEntropyLoss() feature_encoder_optim = torch.optim.Adam(feature_encoder.parameters(), lr=LEARNING_RATE) feature_encoder_scheduler = StepLR(feature_encoder_optim, step_size=100000, gamma=0.5) if os.path.exists(str("./models/omniglot_feature_encoder_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl")): feature_encoder.load_state_dict(torch.load(str("./models/omniglot_feature_encoder_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl"))) print("load feature encoder success") ''' if os.path.exists(str("./models/omniglot_actor_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl")): actor.load_state_dict(torch.load(str("./models/omniglot_actor_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl"))) print("load actor network success") if os.path.exists(str("./models/omniglot_critic_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl")): critic.load_state_dict(torch.load(str("./models/omniglot_critic_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl"))) print("load critic network success") ''' # * Step 3: build graph print("Training...") last_accuracy = 0.0 # embedding_loss_list = [] policy_loss_list = [] value_loss_list = [] for episode in range(EPISODE): # * init dataset # * sample_dataloader is to obtain previous samples for compare # * batch_dataloader is to batch samples for training degrees = random.choice([0, 90, 180, 270]) number_of_query_image = 10 task = tg.OmniglotTask(metatrain_character_folders, CLASS_NUM, SAMPLE_NUM_PER_CLASS, number_of_query_image) sample_dataloader = tg.get_data_loader(task, num_per_class=SAMPLE_NUM_PER_CLASS, split="train", shuffle=False, rotation=degrees) batch_dataloader = tg.get_data_loader(task, num_per_class=number_of_query_image, split="test", shuffle=True, rotation=degrees) # * sample datas # samples, sample_labels = sample_dataloader.__iter__().next() # batches, batch_labels = batch_dataloader.__iter__().next() samples, sample_labels = next(iter(sample_dataloader)) batches, batch_labels = next(iter(batch_dataloader)) # RFT_samples, RFT_sample_labels = samples, sample_labels samples, sample_labels = samples.to(device), sample_labels.to(device) batches, batch_labels = batches.to(device), batch_labels.to(device) # one_hot_sample_labels = torch.zeros(SAMPLE_NUM_PER_CLASS * CLASS_NUM, CLASS_NUM).to(device).scatter_(1, sample_labels.view(-1, 1), 1) # * calculates features sample_features = feature_encoder(samples) # RFT_sample_features = sample_features.detach().cpu().reshape(RFT_samples.shape[0], -1) sample_features = sample_features.view(CLASS_NUM, SAMPLE_NUM_PER_CLASS, FEATURE_DIM, 5, 5) sample_features = torch.sum(sample_features, 1).squeeze(1) batch_features = feature_encoder(batches) # RFT_batch_features = batch_features.detach().cpu().reshape(RFT_batches.shape[0], -1) # embedding_loss = mse(linear, one_hot_sample_labels) # embedding_loss = cross_entropy(linear, sample_labels) # * calculate relations # * each batch sample link to every samples to calculate relations # * to form a 100 * 128 matrix for relation network sample_features_ext = sample_features.unsqueeze(0).repeat(number_of_query_image * CLASS_NUM, 1, 1, 1, 1) batch_features_ext = batch_features.unsqueeze(0).repeat(CLASS_NUM, 1, 1, 1, 1) batch_features_ext = torch.transpose(batch_features_ext, 0, 1) relation_pairs = torch.cat((sample_features_ext, batch_features_ext), 2).view(-1, FEATURE_DIM * 2, 5, 5) policy_loss, value_loss = agent.update_model(relation_pairs, batch_labels) if (episode + 1) % 100 == 0: print(f"episode : {episode+1}, policy_loss : {policy_loss.cpu().detach().numpy()}") print(f"episode : {episode+1}, policy_loss : {value_loss.cpu().detach().numpy()}") policy_loss_list.append(policy_loss.cpu().detach().numpy()) value_loss_list.append(value_loss.cpu().detach().numpy()) if (episode + 1) % 1000 == 0: print("Testing...") total_reward = 0 # feature_encoder.eval() # relation_network.eval() for i in range(TEST_EPISODE): degrees = random.choice([0, 90, 180, 270]) number_of_query_image = 10 task = tg.OmniglotTask(metatest_character_folders, CLASS_NUM, SAMPLE_NUM_PER_CLASS, number_of_query_image) sample_dataloader = tg.get_data_loader(task, num_per_class=SAMPLE_NUM_PER_CLASS, split="train", shuffle=False, rotation=degrees) test_dataloader = tg.get_data_loader(task, num_per_class=number_of_query_image, split="test", shuffle=True, rotation=degrees) sample_images, sample_labels = next(iter(sample_dataloader)) test_images, test_labels = next(iter(test_dataloader)) sample_images, sample_labels = sample_images.to(device), sample_labels.to(device) test_images, test_labels = test_images.to(device), test_labels.to(device) # * calculate features sample_features = feature_encoder(sample_images) sample_features = sample_features.view(CLASS_NUM, SAMPLE_NUM_PER_CLASS, FEATURE_DIM, 5, 5) sample_features = torch.sum(sample_features, 1).squeeze(1) test_features = feature_encoder(test_images) # * calculate relations # * each batch sample link to every samples to calculate relations # * to form a 100x128 matrix for relation network sample_features_ext = sample_features.unsqueeze(0).repeat(number_of_query_image * CLASS_NUM, 1, 1, 1, 1) test_features_ext = test_features.unsqueeze(0).repeat(CLASS_NUM, 1, 1, 1, 1) test_features_ext = torch.transpose(test_features_ext, 0, 1) relation_pairs = torch.cat((sample_features_ext, test_features_ext), 2).view(-1, FEATURE_DIM * 2, 5, 5) rewards = agent.step(relation_pairs, test_labels) total_reward += np.sum(rewards) # feature_encoder.train() # relation_network.train() test_accuracy = total_reward / (1.0 * CLASS_NUM * SAMPLE_NUM_PER_CLASS * TEST_EPISODE) mean_policy_loss = np.mean(policy_loss_list) mean_value_loss = np.mean(value_loss_list) print(f'mean_policy_loss : {mean_policy_loss}') print(f'mean_value_loss : {mean_value_loss}') writer.add_scalar('1.mean_policy_loss', mean_policy_loss, episode + 1) writer.add_scalar('2.mean_value_loss', mean_value_loss, episode + 1) writer.add_scalar('test accuracy', test_accuracy, episode + 1) policy_loss_list = [] value_loss_list = [] if test_accuracy > last_accuracy: # save networks torch.save( feature_encoder.state_dict(), str("./models/omniglot_feature_encoder_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl") ) torch.save( relation_network.state_dict(), str("./models/omniglot_relation_network_" + str(CLASS_NUM) + "way_" + str(SAMPLE_NUM_PER_CLASS) + "shot.pkl") ) print("save networks for episode:", episode) last_accuracy = test_accuracy