├── Morph_2_new ├── results │ ├── confs.npy │ ├── test_losses.npy │ ├── train_losses.npy │ ├── test_accuracies.npy │ ├── train_accuracies.npy │ ├── ground_truth_200_epoch.npy │ └── reconstruction_200_epoch.npy └── code │ └── CapsNet.py ├── Morph_Reger ├── results │ ├── test_losses.npy │ └── train_losses.npy └── code │ └── CapsNet.py ├── CNN_Baseline_Reger ├── results │ ├── test_losses.npy │ └── train_losses.npy └── code │ ├── cnnExercise.py │ ├── DataLoader.py │ └── train_cnn.py ├── CNN_Baseline_Morph_3_Drop_2fc ├── results │ ├── test_accs.npy │ └── train_accs.npy └── code │ ├── cnnExercise.py │ ├── DataLoader.py │ └── train_cnn.py └── README.md /Morph_2_new/results/confs.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/Morph_2_new/results/confs.npy -------------------------------------------------------------------------------- /Morph_2_new/results/test_losses.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/Morph_2_new/results/test_losses.npy -------------------------------------------------------------------------------- /Morph_Reger/results/test_losses.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/Morph_Reger/results/test_losses.npy -------------------------------------------------------------------------------- /Morph_2_new/results/train_losses.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/Morph_2_new/results/train_losses.npy -------------------------------------------------------------------------------- /Morph_Reger/results/train_losses.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/Morph_Reger/results/train_losses.npy -------------------------------------------------------------------------------- /Morph_2_new/results/test_accuracies.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/Morph_2_new/results/test_accuracies.npy -------------------------------------------------------------------------------- /Morph_2_new/results/train_accuracies.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/Morph_2_new/results/train_accuracies.npy -------------------------------------------------------------------------------- /CNN_Baseline_Reger/results/test_losses.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/CNN_Baseline_Reger/results/test_losses.npy -------------------------------------------------------------------------------- /CNN_Baseline_Reger/results/train_losses.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/CNN_Baseline_Reger/results/train_losses.npy -------------------------------------------------------------------------------- /Morph_2_new/results/ground_truth_200_epoch.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/Morph_2_new/results/ground_truth_200_epoch.npy -------------------------------------------------------------------------------- /Morph_2_new/results/reconstruction_200_epoch.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/Morph_2_new/results/reconstruction_200_epoch.npy -------------------------------------------------------------------------------- /CNN_Baseline_Morph_3_Drop_2fc/results/test_accs.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/CNN_Baseline_Morph_3_Drop_2fc/results/test_accs.npy -------------------------------------------------------------------------------- /CNN_Baseline_Morph_3_Drop_2fc/results/train_accs.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RezaKatebi/Galaxy-Morphology-CapsNet/HEAD/CNN_Baseline_Morph_3_Drop_2fc/results/train_accs.npy -------------------------------------------------------------------------------- /CNN_Baseline_Reger/code/cnnExercise.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Author: Reza Katebi 3 | # This code runs the CNN 4 | # run this code for the training 5 | # ====================================================================== 6 | 7 | from __future__ import absolute_import 8 | from __future__ import division 9 | from __future__ import print_function 10 | 11 | import argparse 12 | from train_cnn import ConvNet 13 | 14 | 15 | # Set parameters for CNNs. 16 | parser = argparse.ArgumentParser('CNN Exercise.') 17 | parser.add_argument('--learning_rate', 18 | type=float, default=0.1, 19 | help='Initial learning rate.') 20 | parser.add_argument('--num_epochs', 21 | type=int, 22 | default=30, 23 | help='Number of epochs to run trainer.') 24 | parser.add_argument('--beta', 25 | type=float, 26 | default=0.1, 27 | help='decay rate of L2 regulization.') 28 | parser.add_argument('--batch_size', 29 | type=int, default=5, 30 | help='Batch size. Must divide evenly into the dataset sizes.') 31 | parser.add_argument('--input_data_dir', 32 | type=str, 33 | default='./data/mnist', 34 | help='Directory to put the training data.') 35 | 36 | FLAGS = None 37 | FLAGS, unparsed = parser.parse_known_args() 38 | 39 | 40 | cnn = ConvNet() 41 | accuracy = cnn.train_and_evaluate(FLAGS) 42 | 43 | # Output accuracy 44 | print(20 * '*' + 'model' + 20 * '*') 45 | print('accuracy is %f' % (accuracy)) 46 | print() 47 | -------------------------------------------------------------------------------- /CNN_Baseline_Morph_3_Drop_2fc/code/cnnExercise.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Author: Reza Katebi 3 | # This code runs the CNN 4 | # run this code for the training 5 | # ====================================================================== 6 | 7 | from __future__ import absolute_import 8 | from __future__ import division 9 | from __future__ import print_function 10 | 11 | import argparse 12 | from train_cnn import ConvNet 13 | 14 | 15 | # Set parameters for CNNs. 16 | parser = argparse.ArgumentParser('CNN Exercise.') 17 | parser.add_argument('--learning_rate', 18 | type=float, default=0.01, 19 | help='Initial learning rate.') 20 | parser.add_argument('--num_epochs', 21 | type=int, 22 | default=200, 23 | help='Number of epochs to run trainer.') 24 | parser.add_argument('--beta', 25 | type=float, 26 | default=0.1, 27 | help='decay rate of L2 regulization.') 28 | parser.add_argument('--batch_size', 29 | type=int, default=20, 30 | help='Batch size. Must divide evenly into the dataset sizes.') 31 | parser.add_argument('--input_data_dir', 32 | type=str, 33 | default='./data/mnist', 34 | help='Directory to put the training data.') 35 | 36 | 37 | FLAGS = None 38 | FLAGS, unparsed = parser.parse_known_args() 39 | 40 | 41 | cnn = ConvNet() 42 | accuracy = cnn.train_and_evaluate(FLAGS) 43 | 44 | # Output accuracy 45 | print(20 * '*' + 'model' + 20 * '*') 46 | print('accuracy is %f' % (accuracy)) 47 | print() 48 | -------------------------------------------------------------------------------- /CNN_Baseline_Reger/code/DataLoader.py: -------------------------------------------------------------------------------- 1 | ############################ 2 | # Author: Reza Katebi 3 | # This code loads the dataset 4 | # and divides it to 80 % training 5 | # and 20 % testing data. It 6 | # also makes a batch loader using 7 | # pytorch's dataloader 8 | ########################### 9 | import numpy as np 10 | import torch 11 | # from torchvision import transforms 12 | from torch.utils.data import DataLoader 13 | from torch.utils.data import TensorDataset 14 | from sklearn.model_selection import train_test_split 15 | 16 | 17 | class DataLoad: 18 | """ 19 | Loads the dataset and divides it to train and 20 | test and makes them Tensordataset and loaders 21 | for loading batches. 22 | """ 23 | 24 | def __init__(self, batch_size): 25 | X = np.load('../data/train_downsample.npy')[:61500] 26 | y = np.load('../data/train_labels_regression.npy')[:61500] 27 | print(X.shape) 28 | X_train, X_test, y_train, y_test = train_test_split( 29 | X, y, test_size=0.2, random_state=42) 30 | 31 | X_train = torch.from_numpy(X_train).float() 32 | y_train = torch.from_numpy(y_train) 33 | 34 | X_test = torch.from_numpy(X_test).float() 35 | y_test = torch.from_numpy(y_test) 36 | 37 | train_dataset = TensorDataset(X_train, y_train) 38 | test_dataset = TensorDataset(X_test, y_test) 39 | 40 | self.train_loader = DataLoader( 41 | train_dataset, batch_size=batch_size, shuffle=True) 42 | self.test_loader = DataLoader( 43 | test_dataset, batch_size=batch_size, shuffle=True) 44 | 45 | 46 | if __name__ == "__main__": 47 | Galaxy = DataLoad(100) 48 | a = Galaxy.train_loader 49 | for batch_id, (data, target) in enumerate(a): 50 | 51 | target = torch.eye(5).index_select(dim=0, index=target) 52 | print(data.shape) 53 | print(target.shape) 54 | -------------------------------------------------------------------------------- /CNN_Baseline_Morph_3_Drop_2fc/code/DataLoader.py: -------------------------------------------------------------------------------- 1 | ############################ 2 | # Author: Reza Katebi 3 | # This code loads the dataset 4 | # and divides it to 80 % training 5 | # and 20 % testing data. It 6 | # also makes a batch loader using 7 | # pytorch's dataloader 8 | ########################### 9 | import numpy as np 10 | import torch 11 | # from torchvision import transforms 12 | from torch.utils.data import DataLoader 13 | from torch.utils.data import TensorDataset 14 | from sklearn.model_selection import train_test_split 15 | 16 | 17 | class DataLoad: 18 | """ 19 | Loads the dataset and divides it to train and 20 | test and makes them Tensordataset and loaders 21 | for loading batches. 22 | """ 23 | 24 | def __init__(self, batch_size): 25 | X = np.load('../data/train_downsample_agreed_0.8.npy') 26 | y = np.load('../data/train_labels_agreed_0.8.npy') 27 | print(X.shape) 28 | X_train, X_test, y_train, y_test = train_test_split( 29 | X, y, test_size=0.2, random_state=42) 30 | 31 | X_train = torch.from_numpy(X_train).float() 32 | y_train = torch.from_numpy(y_train) 33 | 34 | X_test = torch.from_numpy(X_test).float() 35 | y_test = torch.from_numpy(y_test) 36 | 37 | train_dataset = TensorDataset(X_train, y_train) 38 | test_dataset = TensorDataset(X_test, y_test) 39 | 40 | self.train_loader = DataLoader( 41 | train_dataset, batch_size=batch_size, shuffle=True) 42 | self.test_loader = DataLoader( 43 | test_dataset, batch_size=batch_size, shuffle=True) 44 | 45 | 46 | if __name__ == "__main__": 47 | Galaxy = DataLoad(100) 48 | a = Galaxy.train_loader 49 | for batch_id, (data, target) in enumerate(a): 50 | 51 | target = torch.eye(5).index_select(dim=0, index=target) 52 | print(data.shape) 53 | print(target.shape) 54 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Galaxy morphology prediction using capsule networks 2 | We implemented our capsule net based on the code provided in [here](https://github.com/gram-ai/capsule-networks). The paper is submitted to MNRAS and you can find it on arXiv at [arXiv:1809.08377](https://arxiv.org/abs/1809.08377). You can also read a news article about our work [here](https://www.ohio-forum.com/2018/10/katebis-class-project-leads-to-a-more-accurate-way-to-classify-galaxy-images/) 3 | 4 | **Datasets** can be dowloaded from [here](https://catmailohio-my.sharepoint.com/:f:/g/personal/rk726014_ohio_edu/Eq-gLBLOB9hMohSCe0ET0O0Bog7_rdZXPziAgy2x2HA6vg?e=5dTHp9). The files were to large so I had to share them through another way. 5 | ## 1 Required Packages 6 | 1. python 3.6 7 | 2. numpy 8 | 3. scipy 9 | 4. [opencv](https://opencv.org/) 10 | 5. pandas 11 | 6. CUDA 12 | 6. [pytorch](https://pytorch.org/) 13 | 7. [torchnet](https://github.com/pytorch/tnt) 14 | 8. [tqdm](https://github.com/tqdm/tqdm) 15 | 9. sklearn 16 | 17 | ## 2 System Requirement 18 | We used one NVIDIA Tesla P100 GPU unit along with 4 CPUs on Owen cluster at Ohio Supercomputer Center (OSC) with 16Gb of memory. 19 | 20 | ## 3 Data Preprocessing 21 | Read the data preprocessing part in the paper. To crop and downsample the dataset from [Kaggle](https://www.kaggle.com/c/galaxy-zoo-the-galaxy-challenge) simply download the dataset to a folder and use the **data_preprocess.py** provided in **preporcess** folder. The preprocessed datasets are provided here. 22 | 23 | ## 4 Folder Structure 24 | Each folder contains the following structure: 25 | - folder/ 26 | - code/ 27 | - data/ 28 | - results/ 29 | - epochs/ (only for CapsNet) 30 | 31 | ## 5 Classification based on answers to question 1 32 | ### 5.1 Baseline CNN 33 | The **CNN_Baseline_Morph_3_Drop_2fc** folder contains the code for this scenario. The datasets have been provided above. The files are *train_downsample_agreed_0.8.npy* and *train_labels_agreed_0.8.npy*. 34 | 35 | ### 5.2 Capsule network 36 | The Morph_2_new folder contains the code and data is the same as provided in baseline model. You can copy and paste it in the dataset folder here. The **epochs** folder is where the model after each epoch of training will be saved and if you want to restart your training from an specific epoch you can load the model and restart your training. 37 | 38 | ## 6 Regression 39 | ## 6.1 Baseline CNN 40 | The **CNN_Baseline_Reger** folder contains the codes. The datasets have been provided above.The files are *train_downsample.npy* and *train_labels_regression.npy*. 41 | 42 | ## 6.2 Capsule Network 43 | The **Morph_Reger** folder contains the code and the dataset is the same as provided in baseline model. 44 | -------------------------------------------------------------------------------- /CNN_Baseline_Reger/code/train_cnn.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Author: Reza Katebi 3 | # This code contains the Baseline CNN model 4 | # ====================================================================== 5 | 6 | import numpy 7 | import time 8 | import torch 9 | import torch.nn as NN 10 | import torch.optim as optim 11 | import torch.nn.functional as F 12 | from torch.autograd import Variable 13 | from DataLoader import DataLoad 14 | 15 | # Change the shape of the input tensor 16 | 17 | 18 | class ViewOP(torch.nn.Module): 19 | def __init__(self, *shape): 20 | super(ViewOP, self).__init__() 21 | self.shape = shape 22 | 23 | def forward(self, input): 24 | # output = input.view(input.size(0), -1) 25 | # print(output.shape) 26 | return input.view(input.size(0), -1) # (self.shape) 27 | 28 | 29 | ########### Convolutional neural network class ############ 30 | class ConvNet(object): 31 | def __init__(self): 32 | if torch.cuda.is_available(): 33 | self.use_gpu = True 34 | else: 35 | self.use_gpu = False 36 | 37 | def ConvBlock(self, num, layers, model, inp_fl, outp_fl): 38 | for i in range(layers): 39 | if i == 0: 40 | model.add_module('Conv2d_{}_{}'.format(num, i), 41 | NN.Conv2d(inp_fl, outp_fl, kernel_size=3, stride=1, padding=(1, 1))) 42 | model.add_module('BN_{}_{}'.format(num, i), NN.BatchNorm2d(outp_fl)) 43 | model.add_module('ReLU_{}_{}'.format(num, i), NN.ReLU()) 44 | else: 45 | model.add_module('Conv2d_{}_{}'.format(num, i), 46 | NN.Conv2d(outp_fl, outp_fl, kernel_size=3, stride=1, padding=(1, 1))) 47 | model.add_module('BN_{}_{}'.format(num, i), NN.BatchNorm2d(outp_fl)) 48 | model.add_module('ReLU_{}_{}'.format(num, i), NN.ReLU()) 49 | model.add_module('Pooling_{}_{}'.format(num, i), NN.MaxPool2d(2)) 50 | 51 | def augmentation(self, x, max_shift=2): 52 | _, _, height, width = x.size() 53 | 54 | h_shift, w_shift = numpy.random.randint(-max_shift, max_shift + 1, size=2) 55 | source_height_slice = slice(max(0, h_shift), h_shift + height) 56 | source_width_slice = slice(max(0, w_shift), w_shift + width) 57 | target_height_slice = slice(max(0, -h_shift), -h_shift + height) 58 | target_width_slice = slice(max(0, -w_shift), -w_shift + width) 59 | 60 | shifted_image = torch.zeros(*x.size()) 61 | shifted_image[:, :, source_height_slice, source_width_slice] = x[:, 62 | :, target_height_slice, target_width_slice] 63 | return shifted_image.float() 64 | 65 | def CNNModel(self, class_num): 66 | 67 | model = NN.Sequential() 68 | model.add_module('Conv2d_1', NN.Conv2d(3, 512, kernel_size=9, stride=1)) 69 | model.add_module('ReLU_1', NN.ReLU()) 70 | model.add_module('Pooling_1', NN.MaxPool2d(2)) 71 | model.add_module('Conv2d_2', NN.Conv2d(512, 256, kernel_size=5, stride=1)) 72 | model.add_module('ReLU_2', NN.ReLU()) 73 | model.add_module('Pooling_1', NN.MaxPool2d(2)) 74 | 75 | # Flatten 76 | model.add_module('Flatten', ViewOP()) 77 | model.add_module('Linear_1', NN.Linear(200704, 1024)) 78 | model.add_module('ReLU_L_1', NN.ReLU()) 79 | # model.add_module('Dropout_1', NN.Dropout(0.5)) 80 | model.add_module('Linear_3', NN.Linear(1024, class_num)) 81 | 82 | return model 83 | 84 | def train(self, model, optimizer, train_loader, ftype, itype): 85 | loss = 0 86 | for _, train_batch in enumerate(train_loader, 0): 87 | batch_x = train_batch[0] / 255 88 | batch_y = train_batch[1] 89 | 90 | X = Variable(self.augmentation(batch_x).type(ftype), requires_grad=False) 91 | Y = Variable(batch_y.type(ftype), requires_grad=False) 92 | 93 | optimizer.zero_grad() 94 | MSE = NN.MSELoss() 95 | loss_func = MSE(model(X), Y) 96 | loss_func.backward() 97 | optimizer.step() 98 | loss += loss_func.item() * batch_x.size(0) 99 | RMSE = numpy.sqrt(loss / len(train_loader.dataset)) 100 | return RMSE 101 | 102 | # Evaluate the trained model on test set. 103 | def evaluate_model(self, model, test_loader, ftype, itype): 104 | model.eval() 105 | loss = 0 106 | for _, test_batch in enumerate(test_loader, 0): 107 | batch_x = test_batch[0] / 255 108 | batch_y = test_batch[1] 109 | 110 | X = Variable(self.augmentation(batch_x).type(ftype), requires_grad=False) 111 | Y = Variable(batch_y.type(ftype), requires_grad=False) 112 | MSE = NN.MSELoss() 113 | loss_func = MSE(model(X), Y) 114 | 115 | # Move tensor from GPU to CPU. 116 | # if self.use_gpu: 117 | # Y = Y.cpu() 118 | loss += loss_func.item() * batch_x.size(0) 119 | 120 | RMSE = numpy.sqrt(loss / len(test_loader.dataset)) 121 | 122 | return RMSE 123 | 124 | # Entry point for training and evaluation. 125 | def train_and_evaluate(self, FLAGS): 126 | class_num = 37 127 | num_epochs = FLAGS.num_epochs 128 | batch_size = FLAGS.batch_size 129 | # learning_rate = FLAGS.learning_rate 130 | 131 | # Set random number generator seed. 132 | torch.manual_seed(1024) 133 | if self.use_gpu: 134 | torch.cuda.manual_seed_all(1024) 135 | numpy.random.seed(1024) 136 | 137 | model = eval("self.CNNModel")(class_num) 138 | 139 | if self.use_gpu: 140 | ftype = torch.cuda.FloatTensor # float type 141 | itype = torch.cuda.LongTensor # int type 142 | model.cuda() 143 | else: 144 | ftype = torch.FloatTensor # float type 145 | itype = torch.LongTensor # int type 146 | 147 | params = model.parameters() 148 | 149 | #-------------------- Create optimizer --------------------- 150 | optimizer = optim.Adam(params) # , lr=learning_rate) 151 | #--------------------- DataLoader -------------------------- 152 | # Loading the custum dataset using DataLoader 153 | Galaxy = DataLoad(batch_size) 154 | test_losses = [] 155 | train_losses = [] 156 | # Train the model. 157 | for i in range(num_epochs): 158 | print(21 * '*', 'epoch', i+1, 21 * '*') 159 | start_time = time.time() 160 | train_loss = self.train(model, optimizer, Galaxy.train_loader, ftype, itype) 161 | train_losses.append(train_loss) 162 | end_time = time.time() 163 | print('the training took: %d(s)' % (end_time - start_time)) 164 | print("train loss is:", train_loss) 165 | 166 | test_loss = self.evaluate_model(model, Galaxy.test_loader, ftype, itype) 167 | 168 | print("test loss is:", test_loss) 169 | test_losses.append(test_loss) 170 | 171 | numpy.save('../results/test_losses', test_losses) 172 | numpy.save('../results/train_losses', train_losses) 173 | 174 | return self.evaluate_model(model, Galaxy.test_loader, ftype, itype) 175 | -------------------------------------------------------------------------------- /Morph_Reger/code/CapsNet.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.setrecursionlimit(15000) 3 | 4 | import torch 5 | import torch.nn.functional as F 6 | from torch import nn 7 | import numpy as np 8 | from sklearn.model_selection import train_test_split 9 | from sklearn.metrics import r2_score 10 | 11 | BATCH_SIZE = 5 12 | NUM_CLASSES = 37 13 | NUM_EPOCHS = 30 14 | NUM_ROUTING_ITERATIONS = 3 15 | 16 | 17 | def softmax(input, dim=1): 18 | transposed_input = input.transpose(dim, len(input.size()) - 1) 19 | softmaxed_output = F.softmax(transposed_input.contiguous().view(-1, 20 | transposed_input.size(-1)), dim=-1) 21 | return softmaxed_output.view(*transposed_input.size()).transpose(dim, len(input.size()) - 1) 22 | 23 | 24 | def augmentation(x, max_shift=2): 25 | _, _, height, width = x.size() 26 | 27 | h_shift, w_shift = np.random.randint(-max_shift, max_shift + 1, size=2) 28 | source_height_slice = slice(max(0, h_shift), h_shift + height) 29 | source_width_slice = slice(max(0, w_shift), w_shift + width) 30 | target_height_slice = slice(max(0, -h_shift), -h_shift + height) 31 | target_width_slice = slice(max(0, -w_shift), -w_shift + width) 32 | 33 | shifted_image = torch.zeros(*x.size()) 34 | shifted_image[:, :, source_height_slice, source_width_slice] = x[:, 35 | :, target_height_slice, target_width_slice] 36 | return shifted_image.float() 37 | 38 | 39 | class CapsuleLayer(nn.Module): 40 | def __init__(self, num_capsules, num_route_nodes, in_channels, out_channels, kernel_size=None, stride=None, 41 | num_iterations=NUM_ROUTING_ITERATIONS): 42 | super(CapsuleLayer, self).__init__() 43 | 44 | self.num_route_nodes = num_route_nodes 45 | self.num_iterations = num_iterations 46 | 47 | self.num_capsules = num_capsules 48 | 49 | if num_route_nodes != -1: 50 | self.route_weights = nn.Parameter(torch.randn( 51 | num_capsules, num_route_nodes, in_channels, out_channels)) 52 | else: 53 | self.capsules = nn.ModuleList( 54 | [nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=0) for _ in 55 | range(num_capsules)]) 56 | 57 | def squash(self, tensor, dim=-1): 58 | squared_norm = (tensor ** 2).sum(dim=dim, keepdim=True) 59 | scale = squared_norm / (1 + squared_norm) 60 | return scale * tensor / torch.sqrt(squared_norm) 61 | 62 | def forward(self, x): 63 | if self.num_route_nodes != -1: 64 | priors = x[None, :, :, None, :] @ self.route_weights[:, None, :, :, :] 65 | 66 | logits = Variable(torch.zeros(*priors.size())).cuda() 67 | for i in range(self.num_iterations): 68 | probs = softmax(logits, dim=2) 69 | outputs = self.squash((probs * priors).sum(dim=2, keepdim=True)) 70 | 71 | if i != self.num_iterations - 1: 72 | delta_logits = (priors * outputs).sum(dim=-1, keepdim=True) 73 | logits = logits + delta_logits 74 | else: 75 | outputs = [capsule(x).view(x.size(0), -1, 1) for capsule in self.capsules] 76 | outputs = torch.cat(outputs, dim=-1) 77 | outputs = self.squash(outputs) 78 | 79 | return outputs 80 | 81 | 82 | class CapsuleNet(nn.Module): 83 | def __init__(self): 84 | super(CapsuleNet, self).__init__() 85 | 86 | self.conv1 = nn.Conv2d(in_channels=3, out_channels=256, kernel_size=9, stride=1) 87 | self.primary_capsules = CapsuleLayer(num_capsules=8, num_route_nodes=-1, in_channels=256, out_channels=32, 88 | kernel_size=9, stride=2) 89 | self.digit_capsules = CapsuleLayer(num_capsules=NUM_CLASSES, num_route_nodes=32 * 28 * 28, in_channels=8, 90 | out_channels=16) 91 | 92 | self.Linear = nn.Linear(16 * NUM_CLASSES, NUM_CLASSES) 93 | 94 | def forward(self, x, y=None): 95 | x = F.relu(self.conv1(x), inplace=True) 96 | x = self.primary_capsules(x) 97 | x = self.digit_capsules(x).squeeze().transpose(0, 1) 98 | x = (x ** 2).sum(dim=-1) ** 0.5 99 | # x = self.Linear(x.view(x.size(0), -1)) 100 | 101 | return x 102 | 103 | 104 | class CapsuleLoss(nn.Module): 105 | def __init__(self): 106 | super(CapsuleLoss, self).__init__() 107 | self.mse = nn.MSELoss(size_average=True) 108 | 109 | def forward(self, labels, x): 110 | 111 | return self.mse(x, labels) 112 | 113 | 114 | if __name__ == "__main__": 115 | from torch.autograd import Variable 116 | from torch.optim import Adam 117 | from torchnet.engine import Engine 118 | # from torchvision.datasets.mnist import MNIST 119 | from tqdm import tqdm 120 | import torchnet as tnt 121 | 122 | model = CapsuleNet() 123 | # model.load_state_dict(torch.load('epochs/epoch_327.pt')) 124 | model.cuda() 125 | 126 | print("# parameters:", sum(param.numel() for param in model.parameters())) 127 | 128 | optimizer = Adam(model.parameters()) 129 | 130 | engine = Engine() 131 | meter_loss = tnt.meter.AverageValueMeter() 132 | outputs_ = [] 133 | lables_ = [] 134 | 135 | capsule_loss = CapsuleLoss() 136 | 137 | def get_iterator(mode): 138 | X = np.load('../data/train_downsample.npy')[:61500] 139 | y = np.load('../data/train_labels_regression.npy')[:61500] 140 | X_train, X_test, y_train, y_test = train_test_split( 141 | X, y, test_size=0.2, random_state=42) 142 | 143 | if mode: 144 | data = torch.from_numpy(X_train).float() 145 | labels = torch.from_numpy(y_train) 146 | else: 147 | data = torch.from_numpy(X_test).float() 148 | labels = torch.from_numpy(y_test) 149 | 150 | tensor_dataset = tnt.dataset.TensorDataset([data, labels]) 151 | return tensor_dataset.parallel(batch_size=BATCH_SIZE, num_workers=1, shuffle=mode) 152 | 153 | def processor(sample): 154 | data, labels, training = sample 155 | 156 | data = augmentation(data.float() / 255.0) 157 | labels = labels.to(torch.float) 158 | 159 | data = Variable(data).cuda() 160 | labels = Variable(labels).cuda() 161 | 162 | if training: 163 | Output = model(data, labels) 164 | else: 165 | Output = model(data) 166 | 167 | loss = capsule_loss(labels, Output) 168 | 169 | return loss, Output 170 | 171 | def reset_meters(): 172 | meter_loss.reset() 173 | 174 | def on_sample(state): 175 | state['sample'].append(state['train']) 176 | 177 | def on_forward(state): 178 | meter_loss.add(state['loss'].item()) 179 | 180 | def on_start_epoch(state): 181 | reset_meters() 182 | state['iterator'] = tqdm(state['iterator']) 183 | 184 | def on_end_epoch(state): 185 | print('[Epoch %d] Training Loss: %.4f' % ( 186 | state['epoch'], np.sqrt(meter_loss.value()[0]))) 187 | 188 | train_losses.append(np.sqrt(meter_loss.value()[0])) 189 | reset_meters() 190 | 191 | engine.test(processor, get_iterator(False)) 192 | 193 | print('[Epoch %d] Testing Loss: %.4f' % ( 194 | state['epoch'], np.sqrt(meter_loss.value()[0]))) 195 | 196 | torch.save(model.state_dict(), '../epochs/epoch_%d.pt' % state['epoch']) 197 | test_losses.append(np.sqrt(meter_loss.value()[0])) 198 | 199 | # def on_start(state): 200 | # state['epoch'] = 327 201 | # 202 | # engine.hooks['on_start'] = on_start 203 | engine.hooks['on_sample'] = on_sample 204 | engine.hooks['on_forward'] = on_forward 205 | engine.hooks['on_start_epoch'] = on_start_epoch 206 | engine.hooks['on_end_epoch'] = on_end_epoch 207 | 208 | train_losses = [] 209 | test_losses = [] 210 | engine.train(processor, get_iterator(True), maxepoch=NUM_EPOCHS, optimizer=optimizer) 211 | 212 | np.save("../results/train_losses", train_losses) 213 | np.save("../results/test_losses", test_losses) 214 | -------------------------------------------------------------------------------- /CNN_Baseline_Morph_3_Drop_2fc/code/train_cnn.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Author: Reza Katebi 3 | # This code contains the Baseline CNN model 4 | # ====================================================================== 5 | 6 | import numpy 7 | import time 8 | import torch 9 | import torch.nn as NN 10 | import torch.optim as optim 11 | import torch.nn.functional as F 12 | from torch.autograd import Variable 13 | from DataLoader import DataLoad 14 | 15 | # Change the shape of the input tensor 16 | 17 | 18 | class ViewOP(torch.nn.Module): 19 | def __init__(self, *shape): 20 | super(ViewOP, self).__init__() 21 | self.shape = shape 22 | 23 | def forward(self, input): 24 | # output = input.view(input.size(0), -1) 25 | # print(output.shape) 26 | return input.view(input.size(0), -1) # (self.shape) 27 | 28 | 29 | ########### Convolutional neural network class ############ 30 | class ConvNet(object): 31 | def __init__(self): 32 | if torch.cuda.is_available(): 33 | self.use_gpu = True 34 | else: 35 | self.use_gpu = False 36 | 37 | def ConvBlock(self, num, layers, model, inp_fl, outp_fl): 38 | for i in range(layers): 39 | if i == 0: 40 | model.add_module('Conv2d_{}_{}'.format(num, i), 41 | NN.Conv2d(inp_fl, outp_fl, kernel_size=3, stride=1, padding=(1, 1))) 42 | model.add_module('BN_{}_{}'.format(num, i), NN.BatchNorm2d(outp_fl)) 43 | model.add_module('ReLU_{}_{}'.format(num, i), NN.ReLU()) 44 | else: 45 | model.add_module('Conv2d_{}_{}'.format(num, i), 46 | NN.Conv2d(outp_fl, outp_fl, kernel_size=3, stride=1, padding=(1, 1))) 47 | model.add_module('BN_{}_{}'.format(num, i), NN.BatchNorm2d(outp_fl)) 48 | model.add_module('ReLU_{}_{}'.format(num, i), NN.ReLU()) 49 | model.add_module('Pooling_{}_{}'.format(num, i), NN.MaxPool2d(2)) 50 | 51 | def CNNModel(self, class_num): 52 | 53 | model = NN.Sequential() 54 | # Conv blocks 55 | model.add_module('Conv2d_1', NN.Conv2d(3, 512, kernel_size=9, stride=1)) 56 | model.add_module('ReLU_1', NN.ReLU()) 57 | model.add_module('Pooling_1', NN.MaxPool2d(2)) 58 | model.add_module('Conv2d_2', NN.Conv2d(512, 256, kernel_size=5, stride=1)) 59 | model.add_module('ReLU_2', NN.ReLU()) 60 | model.add_module('Pooling_1', NN.MaxPool2d(2)) 61 | 62 | # Flatten 63 | model.add_module('Flatten', ViewOP()) 64 | model.add_module('Linear_1', NN.Linear(200704, 1024)) 65 | model.add_module('ReLU_L_1', NN.ReLU()) 66 | model.add_module('Dropout_1', NN.Dropout(0.5)) 67 | # Fully connected layer 2 68 | model.add_module('Linear_2', NN.Linear(1024, 1024)) 69 | model.add_module('ReLU_L_2', NN.ReLU()) 70 | model.add_module('Dropout_2', NN.Dropout(0.5)) 71 | # Output layer 72 | model.add_module('Linear_3', NN.Linear(1024, class_num)) 73 | model.add_module('LogSoftmax', NN.LogSoftmax(dim=1)) 74 | 75 | return model 76 | 77 | def augmentation(self, x, max_shift=2): 78 | _, _, height, width = x.size() 79 | 80 | h_shift, w_shift = numpy.random.randint(-max_shift, max_shift + 1, size=2) 81 | source_height_slice = slice(max(0, h_shift), h_shift + height) 82 | source_width_slice = slice(max(0, w_shift), w_shift + width) 83 | target_height_slice = slice(max(0, -h_shift), -h_shift + height) 84 | target_width_slice = slice(max(0, -w_shift), -w_shift + width) 85 | 86 | shifted_image = torch.zeros(*x.size()) 87 | shifted_image[:, :, source_height_slice, source_width_slice] = x[:, 88 | :, target_height_slice, target_width_slice] 89 | return shifted_image.float() 90 | 91 | def augmentation_rotate(self, x): 92 | x = x.data.numpy() 93 | a = numpy.random.randint(0, 4) 94 | rot_x = numpy.rot90(x, k=a, axes=(2, 3)).copy() 95 | return torch.tensor(rot_x).float() 96 | 97 | def train(self, model, optimizer, train_loader, ftype, itype): 98 | correct = 0 99 | losses = 0 100 | for _, train_batch in enumerate(train_loader, 0): 101 | batch_x = train_batch[0] / 255 102 | batch_y = train_batch[1] 103 | 104 | X = Variable(self.augmentation(batch_x).type(ftype), requires_grad=False) 105 | Y = Variable(batch_y.type(itype), requires_grad=False) 106 | pred_Y = model(X) 107 | 108 | loss_func = F.nll_loss(pred_Y, Y) 109 | losses += loss_func.item() * batch_x.size(0) 110 | optimizer.zero_grad() 111 | loss_func.backward() 112 | optimizer.step() 113 | 114 | _, idx = torch.max(pred_Y, dim=1) 115 | 116 | # Move tensor from GPU to CPU. 117 | if self.use_gpu: 118 | idx = idx.cpu() 119 | Y = Y.cpu() 120 | 121 | idx = idx.data.numpy() 122 | Y = Y.data.numpy() 123 | correct += numpy.sum(idx == Y) 124 | accuracy = correct / len(train_loader.dataset) 125 | return losses/len(train_loader.dataset), accuracy 126 | # Evaluate the trained model on test set. 127 | 128 | def evaluate_model(self, model, test_loader, ftype, itype): 129 | model.eval() 130 | correct = 0 131 | for _, test_batch in enumerate(test_loader, 0): 132 | batch_x = test_batch[0] / 255 133 | batch_y = test_batch[1] 134 | 135 | X = Variable(self.augmentation(batch_x).type(ftype), requires_grad=False) 136 | Y = Variable(batch_y.type(itype), requires_grad=False) 137 | pred_Y = model(X) 138 | _, idx = torch.max(pred_Y, dim=1) 139 | 140 | # Move tensor from GPU to CPU. 141 | if self.use_gpu: 142 | idx = idx.cpu() 143 | Y = Y.cpu() 144 | 145 | idx = idx.data.numpy() 146 | Y = Y.data.numpy() 147 | correct += numpy.sum(idx == Y) 148 | accuracy = correct / len(test_loader.dataset) 149 | 150 | return accuracy 151 | 152 | # Entry point for training and evaluation. 153 | def train_and_evaluate(self, FLAGS): 154 | class_num = 2 155 | num_epochs = FLAGS.num_epochs 156 | batch_size = FLAGS.batch_size 157 | learning_rate = FLAGS.learning_rate 158 | 159 | # Set random number generator seed. 160 | torch.manual_seed(1024) 161 | if self.use_gpu: 162 | torch.cuda.manual_seed_all(1024) 163 | numpy.random.seed(1024) 164 | 165 | model = eval("self.CNNModel")(class_num) 166 | 167 | if self.use_gpu: 168 | ftype = torch.cuda.FloatTensor # float type 169 | itype = torch.cuda.LongTensor # int type 170 | model.cuda() 171 | else: 172 | ftype = torch.FloatTensor # float type 173 | itype = torch.LongTensor # int type 174 | 175 | params = model.parameters() 176 | 177 | #-------------------- Create optimizer --------------------- 178 | optimizer = optim.Adam(params) 179 | #--------------------- DataLoader -------------------------- 180 | # Loading the custum dataset using DataLoader 181 | Galaxy = DataLoad(batch_size) 182 | test_accs = [] 183 | # Train the model. 184 | train_accs = [] 185 | for i in range(num_epochs): 186 | print(21 * '*', 'epoch', i+1, 21 * '*') 187 | start_time = time.time() 188 | loss_train, accuracy = self.train(model, optimizer, Galaxy.train_loader, ftype, itype) 189 | train_accs.append(accuracy) 190 | end_time = time.time() 191 | print('the training took: %d(s)' % (end_time - start_time)) 192 | print('training accuracy is {} and losss is {}'.format(accuracy, loss_train)) 193 | 194 | test_acc = self.evaluate_model(model, Galaxy.test_loader, ftype, itype) 195 | 196 | print("Accuracy of the trained model on test set is", test_acc) 197 | test_accs.append(test_acc) 198 | 199 | numpy.save('../results/test_accs', test_accs) 200 | numpy.save('../results/train_accs', train_accs) 201 | 202 | return self.evaluate_model(model, Galaxy.test_loader, ftype, itype) 203 | -------------------------------------------------------------------------------- /Morph_2_new/code/CapsNet.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.setrecursionlimit(15000) 3 | 4 | import torch 5 | import torch.nn.functional as F 6 | from torch import nn 7 | import numpy as np 8 | from sklearn.model_selection import train_test_split 9 | 10 | BATCH_SIZE = 20 11 | NUM_CLASSES = 2 12 | NUM_EPOCHS = 200 13 | NUM_ROUTING_ITERATIONS = 3 14 | 15 | 16 | def softmax(input, dim=1): 17 | transposed_input = input.transpose(dim, len(input.size()) - 1) 18 | softmaxed_output = F.softmax(transposed_input.contiguous().view(-1, 19 | transposed_input.size(-1)), dim=-1) 20 | return softmaxed_output.view(*transposed_input.size()).transpose(dim, len(input.size()) - 1) 21 | 22 | 23 | def augmentation(x, max_shift=2): 24 | _, _, height, width = x.size() 25 | 26 | h_shift, w_shift = np.random.randint(-max_shift, max_shift + 1, size=2) 27 | source_height_slice = slice(max(0, h_shift), h_shift + height) 28 | source_width_slice = slice(max(0, w_shift), w_shift + width) 29 | target_height_slice = slice(max(0, -h_shift), -h_shift + height) 30 | target_width_slice = slice(max(0, -w_shift), -w_shift + width) 31 | 32 | shifted_image = torch.zeros(*x.size()) 33 | shifted_image[:, :, source_height_slice, source_width_slice] = x[:, 34 | :, target_height_slice, target_width_slice] 35 | return shifted_image.float() 36 | 37 | 38 | class CapsuleLayer(nn.Module): 39 | def __init__(self, num_capsules, num_route_nodes, in_channels, out_channels, kernel_size=None, stride=None, 40 | num_iterations=NUM_ROUTING_ITERATIONS): 41 | super(CapsuleLayer, self).__init__() 42 | 43 | self.num_route_nodes = num_route_nodes 44 | self.num_iterations = num_iterations 45 | 46 | self.num_capsules = num_capsules 47 | 48 | if num_route_nodes != -1: 49 | self.route_weights = nn.Parameter(torch.randn( 50 | num_capsules, num_route_nodes, in_channels, out_channels)) 51 | else: 52 | self.capsules = nn.ModuleList( 53 | [nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=0) for _ in 54 | range(num_capsules)]) 55 | 56 | def squash(self, tensor, dim=-1): 57 | squared_norm = (tensor ** 2).sum(dim=dim, keepdim=True) 58 | scale = squared_norm / (1 + squared_norm) 59 | return scale * tensor / torch.sqrt(squared_norm) 60 | 61 | def forward(self, x): 62 | if self.num_route_nodes != -1: 63 | priors = x[None, :, :, None, :] @ self.route_weights[:, None, :, :, :] 64 | 65 | logits = Variable(torch.zeros(*priors.size())).cuda() 66 | for i in range(self.num_iterations): 67 | probs = softmax(logits, dim=2) 68 | outputs = self.squash((probs * priors).sum(dim=2, keepdim=True)) 69 | 70 | if i != self.num_iterations - 1: 71 | delta_logits = (priors * outputs).sum(dim=-1, keepdim=True) 72 | logits = logits + delta_logits 73 | else: 74 | outputs = [capsule(x).view(x.size(0), -1, 1) for capsule in self.capsules] 75 | outputs = torch.cat(outputs, dim=-1) 76 | outputs = self.squash(outputs) 77 | 78 | return outputs 79 | 80 | 81 | class CapsuleNet(nn.Module): 82 | def __init__(self): 83 | super(CapsuleNet, self).__init__() 84 | 85 | self.conv1 = nn.Conv2d(in_channels=3, out_channels=256, kernel_size=9, stride=1) 86 | self.primary_capsules = CapsuleLayer(num_capsules=8, num_route_nodes=-1, in_channels=256, out_channels=32, 87 | kernel_size=9, stride=2) 88 | self.digit_capsules = CapsuleLayer(num_capsules=NUM_CLASSES, num_route_nodes=32 * 28 * 28, in_channels=8, 89 | out_channels=16) 90 | 91 | self.decoder = nn.Sequential( 92 | nn.Linear(16 * NUM_CLASSES, 512), 93 | nn.ReLU(inplace=True), 94 | nn.Linear(512, 1024), 95 | nn.ReLU(inplace=True), 96 | nn.Linear(1024, 3 * 72 * 72), 97 | nn.Sigmoid() 98 | ) 99 | 100 | def forward(self, x, y=None): 101 | x = F.relu(self.conv1(x), inplace=True) 102 | x = self.primary_capsules(x) 103 | x = self.digit_capsules(x).squeeze().transpose(0, 1) 104 | 105 | classes = (x ** 2).sum(dim=-1) ** 0.5 106 | classes = F.softmax(classes, dim=-1) 107 | 108 | if y is None: 109 | # In all batches, get the most active capsule. 110 | _, max_length_indices = classes.max(dim=1) 111 | # y = Variable(torch.sparse.torch.eye(NUM_CLASSES)).cuda( 112 | # ).index_select(dim=0, index=max_length_indices) 113 | y = Variable(torch.eye(NUM_CLASSES)).cuda( 114 | ).index_select(dim=0, index=max_length_indices) 115 | 116 | reconstructions = self.decoder((x * y[:, :, None]).view(x.size(0), -1)) 117 | 118 | return classes, reconstructions 119 | 120 | 121 | class CapsuleLoss(nn.Module): 122 | def __init__(self): 123 | super(CapsuleLoss, self).__init__() 124 | self.reconstruction_loss = nn.MSELoss(size_average=False) 125 | 126 | def forward(self, images, labels, classes, reconstructions): 127 | left = F.relu(0.9 - classes, inplace=True) ** 2 128 | right = F.relu(classes - 0.1, inplace=True) ** 2 129 | 130 | margin_loss = labels * left + 0.5 * (1. - labels) * right 131 | margin_loss = margin_loss.sum() 132 | 133 | assert torch.numel(images) == torch.numel(reconstructions) 134 | images = images.view(reconstructions.size()[0], -1) 135 | reconstruction_loss = self.reconstruction_loss(reconstructions, images) 136 | 137 | return (margin_loss + 0.0005 * reconstruction_loss) / images.size(0) 138 | 139 | 140 | if __name__ == "__main__": 141 | from torch.autograd import Variable 142 | from torch.optim import Adam 143 | from torchnet.engine import Engine 144 | # from torchvision.datasets.mnist import MNIST 145 | from tqdm import tqdm 146 | import torchnet as tnt 147 | 148 | model = CapsuleNet() 149 | # model.load_state_dict(torch.load('epochs/epoch_327.pt')) 150 | model.cuda() 151 | 152 | print("# parameters:", sum(param.numel() for param in model.parameters())) 153 | 154 | optimizer = Adam(model.parameters()) 155 | 156 | engine = Engine() 157 | meter_loss = tnt.meter.AverageValueMeter() 158 | meter_accuracy = tnt.meter.ClassErrorMeter(accuracy=True) 159 | confusion_meter = tnt.meter.ConfusionMeter(NUM_CLASSES, normalized=True) 160 | 161 | capsule_loss = CapsuleLoss() 162 | 163 | def get_iterator(mode): 164 | X = np.load('../data/train_downsample_agreed_0.8.npy') 165 | y = np.load('../data/train_labels_agreed_0.8.npy') 166 | X_train, X_test, y_train, y_test = train_test_split( 167 | X, y, test_size=0.2, random_state=42) 168 | 169 | if mode: 170 | data = torch.from_numpy(X_train).float() 171 | labels = torch.from_numpy(y_train) 172 | else: 173 | data = torch.from_numpy(X_test).float() 174 | labels = torch.from_numpy(y_test) 175 | 176 | tensor_dataset = tnt.dataset.TensorDataset([data, labels]) 177 | return tensor_dataset.parallel(batch_size=BATCH_SIZE, num_workers=1, shuffle=mode) 178 | 179 | def processor(sample): 180 | data, labels, training = sample 181 | 182 | data = augmentation(data.float() / 255.0) 183 | labels = torch.LongTensor(labels) 184 | 185 | # labels = torch.sparse.torch.eye(NUM_CLASSES).index_select(dim=0, index=labels) 186 | labels = torch.eye(NUM_CLASSES).index_select(dim=0, index=labels) 187 | 188 | data = Variable(data).cuda() 189 | labels = Variable(labels).cuda() 190 | 191 | if training: 192 | classes, reconstructions = model(data, labels) 193 | else: 194 | classes, reconstructions = model(data) 195 | 196 | loss = capsule_loss(data, labels, classes, reconstructions) 197 | 198 | return loss, classes 199 | 200 | def reset_meters(): 201 | meter_accuracy.reset() 202 | meter_loss.reset() 203 | confusion_meter.reset() 204 | 205 | def on_sample(state): 206 | state['sample'].append(state['train']) 207 | 208 | def on_forward(state): 209 | meter_accuracy.add(state['output'].data, torch.LongTensor(state['sample'][1])) 210 | confusion_meter.add(state['output'].data, torch.LongTensor(state['sample'][1])) 211 | meter_loss.add(state['loss'].item()) #.data[0]) 212 | 213 | def on_start_epoch(state): 214 | reset_meters() 215 | state['iterator'] = tqdm(state['iterator']) 216 | 217 | def on_end_epoch(state): 218 | print('[Epoch %d] Training Loss: %.4f (Accuracy: %.2f%%)' % ( 219 | state['epoch'], meter_loss.value()[0], meter_accuracy.value()[0])) 220 | 221 | train_accs.append(meter_accuracy.value()[0]) 222 | train_losses.append(meter_loss.value()[0]) 223 | 224 | reset_meters() 225 | 226 | engine.test(processor, get_iterator(False)) 227 | 228 | print('[Epoch %d] Testing Loss: %.4f (Accuracy: %.2f%%)' % ( 229 | state['epoch'], meter_loss.value()[0], meter_accuracy.value()[0])) 230 | 231 | torch.save(model.state_dict(), '../epochs/epoch_%d.pt' % state['epoch']) 232 | test_accs.append(meter_accuracy.value()[0]) 233 | test_losses.append(meter_loss.value()[0]) 234 | 235 | if state['epoch'] == NUM_EPOCHS: 236 | confs.append(confusion_meter.value()) 237 | 238 | # Reconstruction visualization. 239 | 240 | test_sample = next(iter(get_iterator(False))) 241 | 242 | ground_truth = (test_sample[0].float() / 255.0) 243 | _, reconstructions = model(Variable(ground_truth).cuda()) 244 | reconstruction = reconstructions.cpu().view_as(ground_truth).data 245 | if state['epoch'] % 10 == 0: 246 | np.save('../results/ground_truth_{}_epoch'.format(state['epoch']), ground_truth) 247 | np.save('../results/reconstruction_{}_epoch'.format(state['epoch']), reconstruction) 248 | 249 | # def on_start(state): 250 | # state['epoch'] = 327 251 | # 252 | # engine.hooks['on_start'] = on_start 253 | 254 | engine.hooks['on_sample'] = on_sample 255 | engine.hooks['on_forward'] = on_forward 256 | engine.hooks['on_start_epoch'] = on_start_epoch 257 | engine.hooks['on_end_epoch'] = on_end_epoch 258 | 259 | train_accs = [] 260 | train_losses = [] 261 | test_accs = [] 262 | test_losses = [] 263 | confs = [] 264 | engine.train(processor, get_iterator(True), maxepoch=NUM_EPOCHS, optimizer=optimizer) 265 | 266 | np.save("../results/train_accuracies", train_accs) 267 | np.save("../results/test_accuracies", test_accs) 268 | 269 | np.save("../results/train_losses", train_losses) 270 | np.save("../results/test_losses", test_losses) 271 | 272 | np.save("../results/confs", confs) 273 | --------------------------------------------------------------------------------