├── __init__.py
├── models
    ├── __init__.py
    ├── mvcnn.py
    └── resnet.py
├── LICENSE
├── util.py
├── README.md
├── .gitignore
├── custom_dataset.py
├── logger.py
└── controller.py


/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/models/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2018 René Birkeland
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/util.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | import os
 3 | 
 4 | def logEpoch(logger, model, epoch, loss, accuracy):
 5 |     # 1. Log scalar values (scalar summary)
 6 |     info = {'loss': loss.item(), 'accuracy': accuracy.item()}
 7 | 
 8 |     for tag, value in info.items():
 9 |         logger.scalar_summary(tag, value, epoch)
10 | 
11 |     # 2. Log values and gradients of the parameters (histogram summary)
12 |     for tag, value in model.named_parameters():
13 |         tag = tag.replace('.', '/')
14 |         logger.histo_summary(tag, value.data.cpu().numpy(), epoch)
15 |         logger.histo_summary(tag + '/grad', value.grad.data.cpu().numpy(), epoch)
16 | 
17 |     # 3. Log training images (image summary)
18 |     #info = {'images': images.view(-1, 28, 28)[:10].cpu().numpy()}
19 | 
20 |     #for tag, images in info.items():
21 |         #logger.image_summary(tag, images, epoch)
22 | 
23 | def save_checkpoint(state, model, resnet=None, checkpoint='checkpoint', filename='checkpoint.pth.tar'):
24 |     if resnet:
25 |         filepath = os.path.join(checkpoint, model + str(resnet) + '_' + filename)
26 |     else:
27 |         filepath = os.path.join(checkpoint, model + '_' + filename)
28 |     torch.save(state, filepath)
29 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # MVCNN-PyTorch
 2 | ## Multi-View CNN built on ResNet/AlexNet to classify 3D objects
 3 | A PyTorch implementation of MVCNN using ResNet, inspired by the paper by [Hang Su](http://vis-www.cs.umass.edu/mvcnn/docs/su15mvcnn.pdf).
 4 | MVCNN uses multiple 2D images of 3D objects to classify them. You can use the provided dataset or create your own.
 5 | 
 6 | Also check out my [RotationNet](https://github.com/RBirkeland/RotationNet) implementation whitch outperforms MVCNN (Under construction).
 7 | 
 8 | ![MVCNN](https://preview.ibb.co/eKcJHy/687474703a2f2f7669732d7777772e63732e756d6173732e6564752f6d76636e6e2f696d616765732f6d76636e6e2e706e67.png)
 9 | 
10 | ### Dependencies
11 | * torch
12 | * torchvision
13 | * numpy
14 | * tensorflow (for logging)
15 | 
16 | ### Dataset
17 | ModelNet40 12-view PNG dataset can be downloaded from [Google Drive](https://drive.google.com/file/d/0B4v2jR3WsindMUE3N2xiLVpyLW8/view).
18 | 
19 | You can also create your own 2D dataset from 3D objects (.obj, .stl, and .off), using [BlenderPhong](https://github.com/WeiTang114/BlenderPhong)
20 | 
21 | ### Setup
22 | ```bash
23 | mkdir checkpoint
24 | mkdir logs
25 | ```
26 | 
27 | ### Train
28 | To start training, simply point to the path of the downloaded dataset. All the other settings are optional.
29 | 
30 | ```
31 | python controller.py <path to dataset>  [--depth N] [--model MODEL] [--epochs N] [-b N]
32 |                                         [--lr LR] [--momentum M] [--lr-decay-freq W]
33 |                                         [--lr-decay W] [--print-freq N] [-r PATH] [--pretrained]
34 | ```
35 | 
36 | To resume from a checkpoint, use the -r tag together with the path to the checkpoint file.
37 | 
38 | ### Tensorboard
39 | To view training logs
40 | ```
41 | tensorboard --logdir='logs' --port=6006
42 | ```
43 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
  1 | # Custom
  2 | *.png
  3 | *.tar
  4 | 
  5 | # Byte-compiled / optimized / DLL files
  6 | __pycache__/
  7 | *.py[cod]
  8 | *$py.class
  9 | 
 10 | # C extensions
 11 | *.so
 12 | 
 13 | # Distribution / packaging
 14 | .Python
 15 | build/
 16 | develop-eggs/
 17 | dist/
 18 | downloads/
 19 | eggs/
 20 | .eggs/
 21 | lib/
 22 | lib64/
 23 | parts/
 24 | sdist/
 25 | var/
 26 | wheels/
 27 | *.egg-info/
 28 | .installed.cfg
 29 | *.egg
 30 | MANIFEST
 31 | 
 32 | # PyInstaller
 33 | #  Usually these files are written by a python script from a template
 34 | #  before PyInstaller builds the exe, so as to inject date/other infos into it.
 35 | *.manifest
 36 | *.spec
 37 | 
 38 | # Installer logs
 39 | pip-log.txt
 40 | pip-delete-this-directory.txt
 41 | 
 42 | # Unit test / coverage reports
 43 | htmlcov/
 44 | .tox/
 45 | .coverage
 46 | .coverage.*
 47 | .cache
 48 | nosetests.xml
 49 | coverage.xml
 50 | *.cover
 51 | .hypothesis/
 52 | .pytest_cache/
 53 | 
 54 | # Translations
 55 | *.mo
 56 | *.pot
 57 | 
 58 | # Django stuff:
 59 | *.log
 60 | local_settings.py
 61 | db.sqlite3
 62 | 
 63 | # Flask stuff:
 64 | instance/
 65 | .webassets-cache
 66 | 
 67 | # Scrapy stuff:
 68 | .scrapy
 69 | 
 70 | # Sphinx documentation
 71 | docs/_build/
 72 | 
 73 | # PyBuilder
 74 | target/
 75 | 
 76 | # Jupyter Notebook
 77 | .ipynb_checkpoints
 78 | 
 79 | # pyenv
 80 | .python-version
 81 | 
 82 | # celery beat schedule file
 83 | celerybeat-schedule
 84 | 
 85 | # SageMath parsed files
 86 | *.sage.py
 87 | 
 88 | # Environments
 89 | .env
 90 | .venv
 91 | env/
 92 | venv/
 93 | ENV/
 94 | env.bak/
 95 | venv.bak/
 96 | 
 97 | # Spyder project settings
 98 | .spyderproject
 99 | .spyproject
100 | 
101 | # Rope project settings
102 | .ropeproject
103 | 
104 | # mkdocs documentation
105 | /site
106 | 
107 | # mypy
108 | .mypy_cache/
109 | 


--------------------------------------------------------------------------------
/custom_dataset.py:
--------------------------------------------------------------------------------
 1 | from torch.utils.data.dataset import Dataset
 2 | import os
 3 | from PIL import Image
 4 | 
 5 | class MultiViewDataSet(Dataset):
 6 | 
 7 |     def find_classes(self, dir):
 8 |         classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
 9 |         classes.sort()
10 |         class_to_idx = {classes[i]: i for i in range(len(classes))}
11 | 
12 |         return classes, class_to_idx
13 | 
14 |     def __init__(self, root, data_type, transform=None, target_transform=None):
15 |         self.x = []
16 |         self.y = []
17 |         self.root = root
18 | 
19 |         self.classes, self.class_to_idx = self.find_classes(root)
20 | 
21 |         self.transform = transform
22 |         self.target_transform = target_transform
23 | 
24 |         # root / <label>  / <train/test> / <item> / <view>.png
25 |         for label in os.listdir(root): # Label
26 |             for item in os.listdir(root + '/' + label + '/' + data_type):
27 |                 views = []
28 |                 for view in os.listdir(root + '/' + label + '/' + data_type + '/' + item):
29 |                     views.append(root + '/' + label + '/' + data_type + '/' + item + '/' + view)
30 | 
31 |                 self.x.append(views)
32 |                 self.y.append(self.class_to_idx[label])
33 | 
34 |     # Override to give PyTorch access to any image on the dataset
35 |     def __getitem__(self, index):
36 |         orginal_views = self.x[index]
37 |         views = []
38 | 
39 |         for view in orginal_views:
40 |             im = Image.open(view)
41 |             im = im.convert('RGB')
42 |             if self.transform is not None:
43 |                 im = self.transform(im)
44 |             views.append(im)
45 | 
46 |         return views, self.y[index]
47 | 
48 |     # Override to give PyTorch size of dataset
49 |     def __len__(self):
50 |         return len(self.x)
51 | 


--------------------------------------------------------------------------------
/logger.py:
--------------------------------------------------------------------------------
 1 | import tensorflow as tf
 2 | import numpy as np
 3 | import scipy.misc
 4 | 
 5 | try:
 6 |     from StringIO import StringIO  # Python 2.7
 7 | except ImportError:
 8 |     from io import BytesIO  # Python 3.x
 9 | 
10 | 
11 | class Logger(object):
12 | 
13 |     def __init__(self, log_dir):
14 |         """Create a summary writer logging to log_dir."""
15 |         self.writer = tf.summary.FileWriter(log_dir)
16 | 
17 |     def scalar_summary(self, tag, value, step):
18 |         """Log a scalar variable."""
19 |         summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)])
20 |         self.writer.add_summary(summary, step)
21 | 
22 |     def image_summary(self, tag, images, step):
23 |         """Log a list of images."""
24 | 
25 |         img_summaries = []
26 |         for i, img in enumerate(images):
27 |             # Write the image to a string
28 |             try:
29 |                 s = StringIO()
30 |             except:
31 |                 s = BytesIO()
32 |             scipy.misc.toimage(img).save(s, format="png")
33 | 
34 |             # Create an Image object
35 |             img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(),
36 |                                        height=img.shape[0],
37 |                                        width=img.shape[1])
38 |             # Create a Summary value
39 |             img_summaries.append(tf.Summary.Value(tag='%s/%d' % (tag, i), image=img_sum))
40 | 
41 |         # Create and write Summary
42 |         summary = tf.Summary(value=img_summaries)
43 |         self.writer.add_summary(summary, step)
44 | 
45 |     def histo_summary(self, tag, values, step, bins=1000):
46 |         """Log a histogram of the tensor of values."""
47 | 
48 |         # Create a histogram using numpy
49 |         counts, bin_edges = np.histogram(values, bins=bins)
50 | 
51 |         # Fill the fields of the histogram proto
52 |         hist = tf.HistogramProto()
53 |         hist.min = float(np.min(values))
54 |         hist.max = float(np.max(values))
55 |         hist.num = int(np.prod(values.shape))
56 |         hist.sum = float(np.sum(values))
57 |         hist.sum_squares = float(np.sum(values ** 2))
58 | 
59 |         # Drop the start of the first bin
60 |         bin_edges = bin_edges[1:]
61 | 
62 |         # Add bin edges and counts
63 |         for edge in bin_edges:
64 |             hist.bucket_limit.append(edge)
65 |         for c in counts:
66 |             hist.bucket.append(c)
67 | 
68 |         # Create and write Summary
69 |         summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)])
70 |         self.writer.add_summary(summary, step)
71 |         self.writer.flush()
72 | 


--------------------------------------------------------------------------------
/models/mvcnn.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | import torch.nn as nn
 3 | import torch.utils.model_zoo as model_zoo
 4 | 
 5 | 
 6 | __all__ = ['MVCNN', 'mvcnn']
 7 | 
 8 | 
 9 | model_urls = {
10 |     'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth',
11 | }
12 | 
13 | 
14 | class MVCNN(nn.Module):
15 | 
16 |     def __init__(self, num_classes=1000):
17 |         super(MVCNN, self).__init__()
18 |         self.features = nn.Sequential(
19 |             nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
20 |             nn.ReLU(inplace=True),
21 |             nn.MaxPool2d(kernel_size=3, stride=2),
22 |             nn.Conv2d(64, 192, kernel_size=5, padding=2),
23 |             nn.ReLU(inplace=True),
24 |             nn.MaxPool2d(kernel_size=3, stride=2),
25 |             nn.Conv2d(192, 384, kernel_size=3, padding=1),
26 |             nn.ReLU(inplace=True),
27 |             nn.Conv2d(384, 256, kernel_size=3, padding=1),
28 |             nn.ReLU(inplace=True),
29 |             nn.Conv2d(256, 256, kernel_size=3, padding=1),
30 |             nn.ReLU(inplace=True),
31 |             nn.MaxPool2d(kernel_size=3, stride=2),
32 |         )
33 |         self.classifier = nn.Sequential(
34 |             nn.Dropout(),
35 |             nn.Linear(256 * 6 * 6, 4096),
36 |             nn.ReLU(inplace=True),
37 |             nn.Dropout(),
38 |             nn.Linear(4096, 4096),
39 |             nn.ReLU(inplace=True),
40 |             nn.Linear(4096, num_classes),
41 |         )
42 | 
43 |     def forward(self, x):
44 |         x = x.transpose(0, 1)
45 |         
46 |         view_pool = []
47 |         
48 |         for v in x:
49 |             v = self.features(v)
50 |             v = v.view(v.size(0), 256 * 6 * 6)
51 |             
52 |             view_pool.append(v)
53 |         
54 |         pooled_view = view_pool[0]
55 |         for i in range(1, len(view_pool)):
56 |             pooled_view = torch.max(pooled_view, view_pool[i])
57 |         
58 |         pooled_view = self.classifier(pooled_view)
59 |         return pooled_view
60 | 
61 | 
62 | def mvcnn(pretrained=False, **kwargs):
63 |     r"""MVCNN model architecture from the
64 |     `"Multi-view Convolutional..." <hhttp://vis-www.cs.umass.edu/mvcnn/docs/su15mvcnn.pdf>`_ paper.
65 |     Args:
66 |         pretrained (bool): If True, returns a model pre-trained on ImageNet
67 |     """
68 |     model = MVCNN(**kwargs)
69 |     if pretrained:
70 |         pretrained_dict = model_zoo.load_url(model_urls['alexnet'])
71 |         model_dict = model.state_dict()
72 |         # 1. filter out unnecessary keys
73 |         pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
74 |         # 2. overwrite entries in the existing state dict
75 |         model_dict.update(pretrained_dict)
76 |         # 3. load the new state dict
77 |         model.load_state_dict(model_dict)
78 |     return model


--------------------------------------------------------------------------------
/controller.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn as nn
  3 | import torch.backends.cudnn as cudnn
  4 | from torch.utils.data import DataLoader
  5 | from torch.autograd import Variable
  6 | 
  7 | import torchvision.transforms as transforms
  8 | 
  9 | import argparse
 10 | import numpy as np
 11 | import time
 12 | import os
 13 | 
 14 | from models.resnet import *
 15 | from models.mvcnn import *
 16 | import util
 17 | from logger import Logger
 18 | from custom_dataset import MultiViewDataSet
 19 | 
 20 | MVCNN = 'mvcnn'
 21 | RESNET = 'resnet'
 22 | MODELS = [RESNET,MVCNN]
 23 | 
 24 | parser = argparse.ArgumentParser(description='MVCNN-PyTorch')
 25 | parser.add_argument('data', metavar='DIR', help='path to dataset')
 26 | parser.add_argument('--depth', choices=[18, 34, 50, 101, 152], type=int, metavar='N', default=18, help='resnet depth (default: resnet18)')
 27 | parser.add_argument('--model', '-m', metavar='MODEL', default=RESNET, choices=MODELS,
 28 |                     help='pretrained model: ' + ' | '.join(MODELS) + ' (default: {})'.format(RESNET))
 29 | parser.add_argument('--epochs', default=100, type=int, metavar='N', help='number of total epochs to run (default: 100)')
 30 | parser.add_argument('-b', '--batch-size', default=4, type=int,
 31 |                     metavar='N', help='mini-batch size (default: 4)')
 32 | parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float,
 33 |                     metavar='LR', help='initial learning rate (default: 0.0001)')
 34 | parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
 35 |                     help='momentum (default: 0.9)')
 36 | parser.add_argument('--lr-decay-freq', default=30, type=float,
 37 |                     metavar='W', help='learning rate decay (default: 30)')
 38 | parser.add_argument('--lr-decay', default=0.1, type=float,
 39 |                     metavar='W', help='learning rate decay (default: 0.1)')
 40 | parser.add_argument('--print-freq', '-p', default=10, type=int,
 41 |                     metavar='N', help='print frequency (default: 10)')
 42 | parser.add_argument('-r', '--resume', default='', type=str, metavar='PATH',
 43 |                     help='path to latest checkpoint (default: none)')
 44 | parser.add_argument('--pretrained', dest='pretrained', action='store_true', help='use pre-trained model')
 45 | 
 46 | args = parser.parse_args()
 47 | 
 48 | print('Loading data')
 49 | 
 50 | transform = transforms.Compose([
 51 |     transforms.CenterCrop(500),
 52 |     transforms.Resize(224),
 53 |     transforms.ToTensor(),
 54 | ])
 55 | 
 56 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 57 | 
 58 | # Load dataset
 59 | dset_train = MultiViewDataSet(args.data, 'train', transform=transform)
 60 | train_loader = DataLoader(dset_train, batch_size=args.batch_size, shuffle=True, num_workers=2)
 61 | 
 62 | dset_val = MultiViewDataSet(args.data, 'test', transform=transform)
 63 | val_loader = DataLoader(dset_val, batch_size=args.batch_size, shuffle=True, num_workers=2)
 64 | 
 65 | classes = dset_train.classes
 66 | print(len(classes), classes)
 67 | 
 68 | if args.model == RESNET:
 69 |     if args.depth == 18:
 70 |         model = resnet18(pretrained=args.pretrained, num_classes=len(classes))
 71 |     elif args.depth == 34:
 72 |         model = resnet34(pretrained=args.pretrained, num_classes=len(classes))
 73 |     elif args.depth == 50:
 74 |         model = resnet50(pretrained=args.pretrained, num_classes=len(classes))
 75 |     elif args.depth == 101:
 76 |         model = resnet101(pretrained=args.pretrained, num_classes=len(classes))
 77 |     elif args.depth == 152:
 78 |         model = resnet152(pretrained=args.pretrained, num_classes=len(classes))
 79 |     else:
 80 |         raise Exception('Specify number of layers for resnet in command line. --resnet N')
 81 |     print('Using ' + args.model + str(args.depth))
 82 | else:
 83 |     model = mvcnn(pretrained=args.pretrained,num_classes=len(classes))
 84 |     print('Using ' + args.model)
 85 | 
 86 | model.to(device)
 87 | cudnn.benchmark = True
 88 | 
 89 | print('Running on ' + str(device))
 90 | 
 91 | logger = Logger('logs')
 92 | 
 93 | # Loss and Optimizer
 94 | lr = args.lr
 95 | n_epochs = args.epochs
 96 | criterion = nn.CrossEntropyLoss()
 97 | optimizer = torch.optim.Adam(model.parameters(), lr=lr)
 98 | 
 99 | best_acc = 0.0
100 | best_loss = 0.0
101 | start_epoch = 0
102 | 
103 | 
104 | # Helper functions
105 | def load_checkpoint():
106 |     global best_acc, start_epoch
107 |     # Load checkpoint.
108 |     print('\n==> Loading checkpoint..')
109 |     assert os.path.isfile(args.resume), 'Error: no checkpoint file found!'
110 | 
111 |     checkpoint = torch.load(args.resume)
112 |     best_acc = checkpoint['best_acc']
113 |     start_epoch = checkpoint['epoch']
114 |     model.load_state_dict(checkpoint['state_dict'])
115 |     optimizer.load_state_dict(checkpoint['optimizer'])
116 | 
117 | 
118 | def train():
119 |     train_size = len(train_loader)
120 | 
121 |     for i, (inputs, targets) in enumerate(train_loader):
122 |         # Convert from list of 3D to 4D
123 |         inputs = np.stack(inputs, axis=1)
124 | 
125 |         inputs = torch.from_numpy(inputs)
126 | 
127 |         inputs, targets = inputs.cuda(device), targets.cuda(device)
128 |         inputs, targets = Variable(inputs), Variable(targets)
129 | 
130 |         # compute output
131 |         outputs = model(inputs)
132 |         loss = criterion(outputs, targets)
133 | 
134 |         # compute gradient and do SGD step
135 |         optimizer.zero_grad()
136 |         loss.backward()
137 |         optimizer.step()
138 | 
139 |         if (i + 1) % args.print_freq == 0:
140 |             print("\tIter [%d/%d] Loss: %.4f" % (i + 1, train_size, loss.item()))
141 | 
142 | 
143 | # Validation and Testing
144 | def eval(data_loader, is_test=False):
145 |     if is_test:
146 |         load_checkpoint()
147 | 
148 |     # Eval
149 |     total = 0.0
150 |     correct = 0.0
151 | 
152 |     total_loss = 0.0
153 |     n = 0
154 | 
155 |     for i, (inputs, targets) in enumerate(data_loader):
156 |         with torch.no_grad():
157 |             # Convert from list of 3D to 4D
158 |             inputs = np.stack(inputs, axis=1)
159 | 
160 |             inputs = torch.from_numpy(inputs)
161 | 
162 |             inputs, targets = inputs.cuda(device), targets.cuda(device)
163 |             inputs, targets = Variable(inputs), Variable(targets)
164 | 
165 |             # compute output
166 |             outputs = model(inputs)
167 |             loss = criterion(outputs, targets)
168 | 
169 |             total_loss += loss
170 |             n += 1
171 | 
172 |             _, predicted = torch.max(outputs.data, 1)
173 |             total += targets.size(0)
174 |             correct += (predicted.cpu() == targets.cpu()).sum()
175 | 
176 |     avg_test_acc = 100 * correct / total
177 |     avg_loss = total_loss / n
178 | 
179 |     return avg_test_acc, avg_loss
180 | 
181 | 
182 | # Training / Eval loop
183 | if args.resume:
184 |     load_checkpoint()
185 | 
186 | for epoch in range(start_epoch, n_epochs):
187 |     print('\n-----------------------------------')
188 |     print('Epoch: [%d/%d]' % (epoch+1, n_epochs))
189 |     start = time.time()
190 | 
191 |     model.train()
192 |     train()
193 |     print('Time taken: %.2f sec.' % (time.time() - start))
194 | 
195 |     model.eval()
196 |     avg_test_acc, avg_loss = eval(val_loader)
197 | 
198 |     print('\nEvaluation:')
199 |     print('\tVal Acc: %.2f - Loss: %.4f' % (avg_test_acc.item(), avg_loss.item()))
200 |     print('\tCurrent best val acc: %.2f' % best_acc)
201 | 
202 |     # Log epoch to tensorboard
203 |     # See log using: tensorboard --logdir='logs' --port=6006
204 |     util.logEpoch(logger, model, epoch + 1, avg_loss, avg_test_acc)
205 | 
206 |     # Save model
207 |     if avg_test_acc > best_acc:
208 |         print('\tSaving checkpoint - Acc: %.2f' % avg_test_acc)
209 |         best_acc = avg_test_acc
210 |         best_loss = avg_loss
211 |         util.save_checkpoint({
212 |             'epoch': epoch + 1,
213 |             'state_dict': model.state_dict(),
214 |             'acc': avg_test_acc,
215 |             'best_acc': best_acc,
216 |             'optimizer': optimizer.state_dict(),
217 |         }, args.model, args.depth)
218 | 
219 |     # Decaying Learning Rate
220 |     if (epoch + 1) % args.lr_decay_freq == 0:
221 |         lr *= args.lr_decay
222 |         optimizer = torch.optim.Adam(model.parameters(), lr=lr)
223 |         print('Learning rate:', lr)
224 | 


--------------------------------------------------------------------------------
/models/resnet.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn as nn
  3 | import math
  4 | import torch.utils.model_zoo as model_zoo
  5 | 
  6 | 
  7 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
  8 |            'resnet152']
  9 | 
 10 | 
 11 | model_urls = {
 12 |     'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
 13 |     'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
 14 |     'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
 15 |     'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
 16 |     'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
 17 | }
 18 | 
 19 | 
 20 | def conv3x3(in_planes, out_planes, stride=1):
 21 |     """3x3 convolution with padding"""
 22 |     return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
 23 |                      padding=1, bias=False)
 24 | 
 25 | 
 26 | class BasicBlock(nn.Module):
 27 |     expansion = 1
 28 | 
 29 |     def __init__(self, inplanes, planes, stride=1, downsample=None):
 30 |         super(BasicBlock, self).__init__()
 31 |         self.conv1 = conv3x3(inplanes, planes, stride)
 32 |         self.bn1 = nn.BatchNorm2d(planes)
 33 |         self.relu = nn.ReLU(inplace=True)
 34 |         self.conv2 = conv3x3(planes, planes)
 35 |         self.bn2 = nn.BatchNorm2d(planes)
 36 |         self.downsample = downsample
 37 |         self.stride = stride
 38 | 
 39 |     def forward(self, x):
 40 |         residual = x
 41 | 
 42 |         out = self.conv1(x)
 43 |         out = self.bn1(out)
 44 |         out = self.relu(out)
 45 | 
 46 |         out = self.conv2(out)
 47 |         out = self.bn2(out)
 48 | 
 49 |         if self.downsample is not None:
 50 |             residual = self.downsample(x)
 51 | 
 52 |         out += residual
 53 |         out = self.relu(out)
 54 | 
 55 |         return out
 56 | 
 57 | 
 58 | class Bottleneck(nn.Module):
 59 |     expansion = 4
 60 | 
 61 |     def __init__(self, inplanes, planes, stride=1, downsample=None):
 62 |         super(Bottleneck, self).__init__()
 63 |         self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
 64 |         self.bn1 = nn.BatchNorm2d(planes)
 65 |         self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
 66 |                                padding=1, bias=False)
 67 |         self.bn2 = nn.BatchNorm2d(planes)
 68 |         self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
 69 |         self.bn3 = nn.BatchNorm2d(planes * 4)
 70 |         self.relu = nn.ReLU(inplace=True)
 71 |         self.downsample = downsample
 72 |         self.stride = stride
 73 | 
 74 |     def forward(self, x):
 75 |         residual = x
 76 | 
 77 |         out = self.conv1(x)
 78 |         out = self.bn1(out)
 79 |         out = self.relu(out)
 80 | 
 81 |         out = self.conv2(out)
 82 |         out = self.bn2(out)
 83 |         out = self.relu(out)
 84 | 
 85 |         out = self.conv3(out)
 86 |         out = self.bn3(out)
 87 | 
 88 |         if self.downsample is not None:
 89 |             residual = self.downsample(x)
 90 | 
 91 |         out += residual
 92 |         out = self.relu(out)
 93 | 
 94 |         return out
 95 | 
 96 | 
 97 | class ResNet(nn.Module):
 98 | 
 99 |     def __init__(self, block, layers, num_classes=3):
100 |         self.inplanes = 64
101 |         super(ResNet, self).__init__()
102 |         self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
103 |         self.bn1 = nn.BatchNorm2d(64)
104 |         self.relu = nn.ReLU(inplace=True)
105 |         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
106 |         self.layer1 = self._make_layer(block, 64, layers[0])
107 |         self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
108 |         self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
109 |         self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
110 |         self.avgpool = nn.AvgPool2d(7, stride=1) #nn.AdaptiveAvgPool2d(7)
111 |         self.fc = nn.Linear(512 * block.expansion, num_classes)
112 | 
113 |         for m in self.modules():
114 |             if isinstance(m, nn.Conv2d):
115 |                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
116 |             elif isinstance(m, nn.BatchNorm2d):
117 |                 nn.init.constant_(m.weight, 1)
118 |                 nn.init.constant_(m.bias, 0)
119 | 
120 |     def _make_layer(self, block, planes, blocks, stride=1):
121 |         downsample = None
122 |         if stride != 1 or self.inplanes != planes * block.expansion:
123 |             downsample = nn.Sequential(
124 |                 nn.Conv2d(self.inplanes, planes * block.expansion,
125 |                           kernel_size=1, stride=stride, bias=False),
126 |                 nn.BatchNorm2d(planes * block.expansion),
127 |             )
128 | 
129 |         layers = []
130 |         layers.append(block(self.inplanes, planes, stride, downsample))
131 |         self.inplanes = planes * block.expansion
132 |         for i in range(1, blocks):
133 |             layers.append(block(self.inplanes, planes))
134 | 
135 |         return nn.Sequential(*layers)
136 | 
137 |     def forward(self, x):
138 |         # Swap batch and views dims
139 |         x = x.transpose(0, 1)
140 |         #print(x.shape)
141 | 
142 |         # View pool
143 |         view_pool = []
144 | 
145 |         for v in x:
146 |             v = self.conv1(v)
147 |             v = self.bn1(v)
148 |             v = self.relu(v)
149 |             v = self.maxpool(v)
150 | 
151 |             v = self.layer1(v)
152 |             v = self.layer2(v)
153 |             v = self.layer3(v)
154 |             v = self.layer4(v)
155 | 
156 |             v = self.avgpool(v)
157 |             v = v.view(v.size(0), -1)
158 | 
159 |             view_pool.append(v)
160 | 
161 |         pooled_view = view_pool[0]
162 |         for i in range(1, len(view_pool)):
163 |             pooled_view = torch.max(pooled_view, view_pool[i])
164 | 
165 |         pooled_view = self.fc(pooled_view)
166 | 
167 |         return pooled_view
168 | 
169 | 
170 | def resnet18(pretrained=False, **kwargs):
171 |     """Constructs a ResNet-18 model.
172 | 
173 |     Args:
174 |         pretrained (bool): If True, returns a model pre-trained on ImageNet
175 |     """
176 |     model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
177 |     if pretrained:
178 |         pretrained_dict = model_zoo.load_url(model_urls['resnet18'])
179 |         model_dict = model.state_dict()
180 |         pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
181 |         model_dict.update(pretrained_dict)
182 |         model.load_state_dict(model_dict)
183 |     return model
184 | 
185 | 
186 | def resnet34(pretrained=False, **kwargs):
187 |     """Constructs a ResNet-34 model.
188 | 
189 |     Args:
190 |         pretrained (bool): If True, returns a model pre-trained on ImageNet
191 |     """
192 |     model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
193 |     if pretrained:
194 |         pretrained_dict = model_zoo.load_url(model_urls['resnet34'])
195 |         model_dict = model.state_dict()
196 |         pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
197 |         model_dict.update(pretrained_dict)
198 |         model.load_state_dict(model_dict)
199 |     return model
200 | 
201 | 
202 | def resnet50(pretrained=False, **kwargs):
203 |     """Constructs a ResNet-50 model.
204 | 
205 |     Args:
206 |         pretrained (bool): If True, returns a model pre-trained on ImageNet
207 |     """
208 |     model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
209 |     if pretrained:
210 |         pretrained_dict = model_zoo.load_url(model_urls['resnet50'])
211 |         model_dict = model.state_dict()
212 |         pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
213 |         model_dict.update(pretrained_dict)
214 |         model.load_state_dict(model_dict)
215 |     return model
216 | 
217 | 
218 | def resnet101(pretrained=False, **kwargs):
219 |     """Constructs a ResNet-101 model.
220 | 
221 |     Args:
222 |         pretrained (bool): If True, returns a model pre-trained on ImageNet
223 |     """
224 |     model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
225 |     if pretrained:
226 |         pretrained_dict = model_zoo.load_url(model_urls['resnet101'])
227 |         model_dict = model.state_dict()
228 |         pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
229 |         model_dict.update(pretrained_dict)
230 |         model.load_state_dict(model_dict)
231 |     return model
232 | 
233 | 
234 | def resnet152(pretrained=False, **kwargs):
235 |     """Constructs a ResNet-152 model.
236 | 
237 |     Args:
238 |         pretrained (bool): If True, returns a model pre-trained on ImageNet
239 |     """
240 |     model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
241 |     if pretrained:
242 |         pretrained_dict = model_zoo.load_url(model_urls['resnet152'])
243 |         model_dict = model.state_dict()
244 |         pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
245 |         model_dict.update(pretrained_dict)
246 |         model.load_state_dict(model_dict)
247 |     return model
248 | 


--------------------------------------------------------------------------------