Skip to content

Commit

Permalink
fix softmax&loss func
Browse files Browse the repository at this point in the history
  • Loading branch information
panthersuper committed Dec 4, 2017
1 parent b728266 commit a6e52b4
Show file tree
Hide file tree
Showing 8 changed files with 225 additions and 148 deletions.
42 changes: 26 additions & 16 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,14 @@
# Dataset Parameters
load_size =256
fine_size = 224
data_mean = np.asarray([0,0,0,0])
batch_size = 8
data_mean = np.asarray([0.3,0.3,0.3,0])
batch_size = 120
voxel_size = 256

# Training Parameters
learning_rate = 0.1
learning_rate = 0.01
training_epoches = 10
step_display = 100
step_display = 10
step_save = 2
path_save = 'recon0'
start_from = ''#'./alexnet64/Epoch28'
Expand All @@ -34,7 +34,9 @@
'fine_size': fine_size,
'voxel_size': voxel_size,
'data_mean': data_mean,
'randomize': True
'randomize': True,
'down_sample_scale':8

}
opt_data_val = {
'img_root': 'data/train_imgs/', # MODIFY PATH ACCORDINGLY
Expand All @@ -43,7 +45,9 @@
'fine_size': fine_size,
'voxel_size': voxel_size,
'data_mean': data_mean,
'randomize': True
'randomize': True,
'down_sample_scale':8

}

def get_accuracy(loader, size, net):
Expand All @@ -69,7 +73,7 @@ def get_accuracy(loader, size, net):

def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
if classname.find('Conv2') != -1:
nn.init.kaiming_uniform(m.weight.data)

loader_train = DataLoaderDisk(**opt_data_train)
Expand All @@ -80,12 +84,14 @@ def weights_init(m):
if start_from != '':
net.load_state_dict(torch.load(start_from))
else:
pass #net.apply(weights_init)
net.apply(weights_init)

# criterion = nn.MSELoss().cuda()
criterion = nn.NLLLoss().cuda()

criterion = nn.MSELoss().cuda()

optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.005)
scheduler = s.StepLR(optimizer, step_size=4, gamma=0.1)
optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.8, weight_decay=0.005)
scheduler = s.StepLR(optimizer, step_size=3, gamma=0.1)

running_loss = 0.0

Expand All @@ -107,7 +113,7 @@ def weights_init(m):
inputs = np.swapaxes(inputs,1,3)
inputs = np.swapaxes(inputs,2,3)
inputs = torch.from_numpy(inputs).float()
labels = torch.from_numpy(labels).float()
labels = torch.from_numpy(labels).long()

# wrap them in Variable
# inputs, labels = Variable(inputs), Variable(labels)
Expand All @@ -118,19 +124,23 @@ def weights_init(m):
optimizer.zero_grad()

# forward + backward + optimize
output = net(inputs) # places output
output = net(inputs).float().contiguous() # places output

output = output.view(batch_size*32*32*32,-1)
labels = labels.view(-1)

loss = criterion(output, labels)

loss = criterion(output, labels).contiguous()
loss.backward()
optimizer.step()

# print statistics
running_loss += loss.data[0]
if i % step_display == step_display - 1: # print every 100 mini-batches
print('PLACES TRAINING Epoch: %d %d loss: %.3f' %
print('PLACES TRAINING Epoch: %d %d loss: %.10f' %
(epoch + starting_num, i + 1, running_loss/step_display))
with open('./' + path_save + '/log.txt', 'a') as f:
f.write('PLACES TRAINING Epoch: %d %d loss: %.3f\n' %
f.write('PLACES TRAINING Epoch: %d %d loss: %.10f\n' %
(epoch + starting_num, i + 1, running_loss/step_display))

running_loss = 0.0
Expand Down
Binary file modified model/__pycache__/reconNet.cpython-36.pyc
Binary file not shown.
157 changes: 130 additions & 27 deletions model/reconNet.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,59 +4,161 @@
sys.path.append('model/')

from alexnet import alex_net
import torch.nn.functional as F

class ReconNet(nn.Module):

def __init__(self):
super(ReconNet, self).__init__()
self.features = alex_net(pretrained=True)
self.features = nn.Sequential(
nn.Conv2d(4, 64, kernel_size=11, stride=4, padding=2),
nn.LeakyReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(64, 192, kernel_size=5, padding=2),
nn.LeakyReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(192, 384, kernel_size=3, padding=1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3, padding=1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.LeakyReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
)

self.latentV = nn.Sequential(
nn.Dropout(),
nn.Linear(256 * 6 * 6, 4096),
nn.BatchNorm1d(4096),
nn.ReLU(inplace=True),
nn.LeakyReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.BatchNorm1d(4096),
nn.ReLU(inplace=True), # latent vector, size:4096
nn.LeakyReLU(inplace=True), # latent vector, size:4096
)

# self.decoding = nn.Sequential(
# # decoding process
# # 2*2*2 de-conv3

# nn.BatchNorm3d(512),
# nn.ReLU(inplace=True),

# nn.ConvTranspose3d(512, 128, kernel_size=4, stride=4, padding=0), #8
# nn.BatchNorm3d(128, affine=True),
# nn.ReLU(inplace=True),

# nn.ConvTranspose3d(128, 32, kernel_size=4, stride=4, padding=0), #32
# nn.BatchNorm3d(32, affine=True),
# nn.ReLU(inplace=True),

# nn.ConvTranspose3d(32, 8, kernel_size=4, stride=4, padding=0), #128
# nn.BatchNorm3d(8, affine=True),
# nn.ReLU(inplace=True),

# nn.ConvTranspose3d(8, 1, kernel_size=2, stride=2, padding=0), #256
# nn.Tanh(),

# )

# self.decoding = nn.Sequential(
# # decoding process
# # 4/4/4 de-conv3

# nn.BatchNorm3d(64),
# nn.ReLU(inplace=True),

# UpsampleConv3Layer(64, 32, kernel_size=3, stride=1, upsample=2, outsize=8), #8
# nn.BatchNorm3d(32, affine=True),
# nn.ReLU(inplace=True),

# UpsampleConv3Layer(32, 8, kernel_size=3, stride=1, upsample=4, outsize=32), #32
# nn.BatchNorm3d(8, affine=True),
# nn.ReLU(inplace=True),

# UpsampleConv3Layer(8, 4, kernel_size=3, stride=1, upsample=4, outsize=128), #128
# nn.BatchNorm3d(4, affine=True),
# nn.ReLU(inplace=True),

# UpsampleConv3Layer(4, 1, kernel_size=3, stride=1, upsample=2, outsize=256), #256
# nn.Tanh(),

# )

# self.decoding = nn.Sequential(
# # decoding process
# # 4/4/4 de-conv3

# nn.BatchNorm3d(64),
# nn.ReLU(inplace=True),

# nn.ConvTranspose3d(64, 32, kernel_size=6, stride=2, padding=2), #8
# nn.BatchNorm3d(32, affine=True),
# nn.ReLU(inplace=True),

# nn.ConvTranspose3d(32, 16, kernel_size=6, stride=2, padding=2), #16
# nn.BatchNorm3d(16, affine=True),
# nn.ReLU(inplace=True),

# nn.ConvTranspose3d(16, 8, kernel_size=6, stride=2, padding=2), #32
# nn.BatchNorm3d(8, affine=True),
# nn.ReLU(inplace=True),

# nn.ConvTranspose3d(8, 4, kernel_size=6, stride=2, padding=2), #64
# nn.BatchNorm3d(4, affine=True),
# nn.ReLU(inplace=True),

# nn.ConvTranspose3d(4, 2, kernel_size=6, stride=2, padding=2), #128
# nn.BatchNorm3d(2, affine=True),
# nn.ReLU(inplace=True),

# nn.ConvTranspose3d(2, 1, kernel_size=6, stride=2, padding=2), #256
# nn.Tanh(),

# )

self.decoding = nn.Sequential(
# decoding process
# 2*2*2 de-conv3

nn.BatchNorm3d(512),
nn.ReLU(inplace=True),
# 4/4/4 de-conv3 -> 32*32*32

nn.ConvTranspose3d(512, 128, kernel_size=4, stride=4, padding=0), #8
nn.BatchNorm3d(128, affine=True),
nn.ReLU(inplace=True),
nn.BatchNorm3d(64),
nn.LeakyReLU(inplace=True),

nn.ConvTranspose3d(128, 32, kernel_size=4, stride=4, padding=0), #32
nn.ConvTranspose3d(64, 32, kernel_size=6, stride=2, padding=2), #8
nn.BatchNorm3d(32, affine=True),
nn.ReLU(inplace=True),
nn.LeakyReLU(inplace=True),

nn.ConvTranspose3d(32, 8, kernel_size=4, stride=4, padding=0), #128
nn.ConvTranspose3d(32, 8, kernel_size=6, stride=2, padding=2), #16
nn.BatchNorm3d(8, affine=True),
nn.ReLU(inplace=True),
nn.LeakyReLU(inplace=True),

nn.ConvTranspose3d(8, 1, kernel_size=2, stride=2, padding=0), #256
nn.ConvTranspose3d(8, 2, kernel_size=6, stride=2, padding=2), #32 60*2*32*32*32
nn.Tanh(),

)

# self.decoding = nn.Sequential(
# # decoding process
# # 4/4/4 de-conv3 -> 32*32*32

# nn.Linear(64 * 4 * 4 * 4, 32768*2),
# nn.Tanh(),

# )

def forward(self, x):
x = self.features(x)
# x = x.view(x.size(0), 256 * 6 * 6)
x = x.view(x.size(0), 256 * 6 * 6)
x = self.latentV(x) # latent vector, size:4096
x = x.view(x.size(0),512,2,2,2) # reshape to 2 by 2 by 2 cube with 512 channels
x = x.view(x.size(0),64,4,4,4) # reshape to 4/4/4 cube with 64 channels
x = self.decoding(x) # convert to 3D voxel distribution

return x

# x = x.view(x.size(0),1,32,32,32)
x = x.view(x.size(0),2,32768) # 60*2*(32*32*32)
x = F.log_softmax(x) # 60*2*(32*32*32)
x = x.permute(0,2,1) # 60*(32*32*32)*2
# return x
return x#.max(0)[1] #flat

class UpsampleConv3Layer(nn.Module):
"""UpsampleConvLayer
Expand All @@ -65,20 +167,21 @@ class UpsampleConv3Layer(nn.Module):
ref: http://distill.pub/2016/deconv-checkerboard/
"""

def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None):
def __init__(self, in_channels, out_channels, kernel_size, stride, outsize, upsample=None):
super(UpsampleConv3Layer, self).__init__()
self.upsample = upsample
if upsample:
self.upsample_layer = nn._functions.thnn.UpsamplingNearest3d(scale_factor=upsample)
self.outsize = outsize
# if upsample:
# self.upsample_layer = nn.functional.upsample(scale_factor=upsample)
reflection_padding = kernel_size // 2
self.reflection_pad = nn.ReplicationPad3d(reflection_padding) #may modify to ReflectionPad
self.conv3d = nn.Conv3d(in_channels, out_channels, kernel_size, stride)
self.conv3d = nn.Conv3d(in_channels, out_channels, kernel_size, stride,padding=reflection_padding)

def forward(self, x):
x_in = x.contiguous()
x_in = x
if self.upsample:
x_in = nn._functions.thnn.UpsamplingNearest3d.apply(x_in, [256,256,256], self.upsample)
x_in = nn.functional.upsample(x_in,scale_factor=self.upsample,size=self.outsize,mode='trilinear')

out = self.reflection_pad(x_in.contiguous())
out = self.conv3d(out)
# out = self.reflection_pad(x_in)
out = self.conv3d(x_in)
return out
39 changes: 39 additions & 0 deletions nohup.out
Original file line number Diff line number Diff line change
Expand Up @@ -465,3 +465,42 @@ Variable containing:
1.00000e-02 *
3.6148
[torch.cuda.FloatTensor of size 1 (GPU 0)]
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1503970438496/work/torch/lib/THC/generic/THCStorage.cu line=66 error=2 : out of memory
# Images found: 404453
# Images found: 404453
Traceback (most recent call last):
File "main.py", line 114, in <module>
inputs, labels= Variable(inputs.cuda()), Variable(labels.cuda())
File "/home/pwz/miniconda3/lib/python3.6/site-packages/torch/_utils.py", line 66, in _cuda
return new_type(self.size()).copy_(self, async)
RuntimeError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1503970438496/work/torch/lib/THC/generic/THCStorage.cu:66
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1503970438496/work/torch/lib/THC/generic/THCStorage.cu line=66 error=2 : out of memory
# Images found: 404453
# Images found: 404453
Traceback (most recent call last):
File "main.py", line 114, in <module>
inputs, labels= Variable(inputs.cuda()), Variable(labels.cuda())
File "/home/pwz/miniconda3/lib/python3.6/site-packages/torch/_utils.py", line 66, in _cuda
return new_type(self.size()).copy_(self, async)
RuntimeError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1503970438496/work/torch/lib/THC/generic/THCStorage.cu:66
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1503970438496/work/torch/lib/THC/generic/THCStorage.cu line=66 error=2 : out of memory
# Images found: 404453
# Images found: 404453
Traceback (most recent call last):
File "main.py", line 121, in <module>
output = net(inputs) # places output
File "/home/pwz/miniconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 224, in __call__
result = self.forward(*input, **kwargs)
File "/home/pwz/sambavol/3dreconstruction/model/reconNet.py", line 113, in forward
x = self.decoding(x) # convert to 3D voxel distribution
File "/home/pwz/miniconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 224, in __call__
result = self.forward(*input, **kwargs)
File "/home/pwz/miniconda3/lib/python3.6/site-packages/torch/nn/modules/container.py", line 67, in forward
input = module(input)
File "/home/pwz/miniconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 224, in __call__
result = self.forward(*input, **kwargs)
File "/home/pwz/miniconda3/lib/python3.6/site-packages/torch/nn/modules/conv.py", line 663, in forward
output_padding, self.groups, self.dilation)
File "/home/pwz/miniconda3/lib/python3.6/site-packages/torch/nn/functional.py", line 203, in conv_transpose3d
return f(input, weight, bias)
RuntimeError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1503970438496/work/torch/lib/THC/generic/THCStorage.cu:66
Loading

0 comments on commit a6e52b4

Please sign in to comment.