Advertisement
Guest User

Untitled

a guest
Jan 28th, 2019
93
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import torch
  2. import torch.nn as nn
  3. import torchvision.transforms as transforms
  4. import torchvision.datasets as dsets
  5. import os
  6. import pandas as pd
  7. import numpy as np
  8. import matplotlib.pyplot as plt
  9. from skimage import io, transform, color
  10. from torch.utils.data import Dataset, DataLoader
  11. from torchvision import transforms, utils
  12. from pathlib import Path
  13. import json
  14.  
  15. '''
  16. ########################################## STEP 1: LOADING DATASET
  17. '''
  18.  
  19. class CarvanaDataset(Dataset):
  20.    
  21.     def __init__(self, img_dir, mask_dir, train):
  22.         """
  23.        Args:
  24.        root_dir (string): Directory with all the images
  25.        train (boolean): Whether the dataset is training data (train = True, test = False)
  26.        """
  27.         self.img_dir = img_dir
  28.         self.mask_dir = mask_dir
  29.         self.train = train
  30.         self.image_list = []
  31.         self.counter = 0 #dummy variable to break out of io loop for faster testing
  32.        
  33.         # Now iterate through all images in the directory to pull them into python
  34.  
  35.         for filename in os.listdir(img_dir):
  36.             img = io.imread(os.path.join(img_dir, filename)) #this joins the paths of my root directory with each filename
  37.             maskname = filename[:-4] + "_mask.png" #cuts off .jpg and adds '_mask.png'
  38.             mask = io.imread(os.path.join(mask_dir, maskname))
  39.             self.image_list.append((torch.tensor(color.rgb2gray(img)), torch.tensor(mask))) #converts car image to grayscale
  40.             self.counter = self.counter + 1
  41.            
  42.             if (self.counter > 100):
  43.                 break
  44.        
  45.     def __len__(self):
  46.         return len(self.image_list)
  47.        
  48.     def __getitem__(self, idx):
  49.        
  50.         return self.image_list[idx]
  51.  
  52. data_set = CarvanaDataset(img_dir = 'train-128/', mask_dir = 'train_masks-128/', train = True)
  53.  
  54. # Split the training data into train and test data
  55. train_dataset = data_set[:int(len(data_set)/2)]
  56. test_dataset = data_set[int(len(data_set)/2):]
  57. plt.imshow(test_dataset[50][0]) #show an image
  58.  
  59. '''
  60. ########################################## STEP 2: MAKING DATASET ITERABLE
  61. '''
  62.  
  63. batch_size = 5 #we will feed the model 100 images at a time
  64. n_iters = 300
  65. num_epochs = n_iters / (len(train_dataset) / batch_size) #need to review epochs and why this is the way to compute it
  66. num_epochs = int(num_epochs)
  67.  
  68. train_loader = torch.utils.data.DataLoader(dataset = train_dataset,
  69.                                            batch_size = batch_size,
  70.                                            shuffle = True) #shuffle ensures we traverse images in different order across epochs
  71.  
  72. test_loader = torch.utils.data.DataLoader(dataset = test_dataset,
  73.                                           batch_size = batch_size,
  74.                                           shuffle = False) #we don't do shuffle here because we only do 1 forward pass
  75.  
  76. '''
  77. ########################################## STEP 3: CREATE MODEL CLASS
  78. '''
  79. class UNet(nn.Module):
  80.     def __init__(self, num_channels=1, num_classes=2):
  81.         super(UNet, self).__init__()
  82.         num_feat = [64, 128, 256, 512, 1024]
  83.  
  84.         self.down1 = nn.Sequential(Conv3x3(num_channels, num_feat[0]))
  85.  
  86.         self.down2 = nn.Sequential(nn.MaxPool2d(kernel_size=2),
  87.                                    Conv3x3(num_feat[0], num_feat[1]))
  88.  
  89.         self.down3 = nn.Sequential(nn.MaxPool2d(kernel_size=2),
  90.                                    Conv3x3(num_feat[1], num_feat[2]))
  91.  
  92.         self.down4 = nn.Sequential(nn.MaxPool2d(kernel_size=2),
  93.                                    Conv3x3(num_feat[2], num_feat[3]))
  94.  
  95.         self.bottom = nn.Sequential(nn.MaxPool2d(kernel_size=2),
  96.                                     Conv3x3(num_feat[3], num_feat[4]))
  97.  
  98.         self.up1 = UpConcat(num_feat[4], num_feat[3])
  99.         self.upconv1 = Conv3x3(num_feat[4], num_feat[3])
  100.  
  101.         self.up2 = UpConcat(num_feat[3], num_feat[2])
  102.         self.upconv2 = Conv3x3(num_feat[3], num_feat[2])
  103.  
  104.         self.up3 = UpConcat(num_feat[2], num_feat[1])
  105.         self.upconv3 = Conv3x3(num_feat[2], num_feat[1])
  106.  
  107.         self.up4 = UpConcat(num_feat[1], num_feat[0])
  108.         self.upconv4 = Conv3x3(num_feat[1], num_feat[0])
  109.  
  110.         self.final = nn.Sequential(nn.Conv2d(num_feat[0],
  111.                                              num_classes,
  112.                                              kernel_size=1),
  113.                                    nn.Softmax2d())
  114.  
  115.     def forward(self, inputs, return_features=False):
  116.         # print(inputs.data.size())
  117.         down1_feat = self.down1(inputs)
  118.         # print(down1_feat.size())
  119.         down2_feat = self.down2(down1_feat)
  120.         # print(down2_feat.size())
  121.         down3_feat = self.down3(down2_feat)
  122.         # print(down3_feat.size())
  123.         down4_feat = self.down4(down3_feat)
  124.         # print(down4_feat.size())
  125.         bottom_feat = self.bottom(down4_feat)
  126.  
  127.         # print(bottom_feat.size())
  128.         up1_feat = self.up1(bottom_feat, down4_feat)
  129.         # print(up1_feat.size())
  130.         up1_feat = self.upconv1(up1_feat)
  131.         # print(up1_feat.size())
  132.         up2_feat = self.up2(up1_feat, down3_feat)
  133.         # print(up2_feat.size())
  134.         up2_feat = self.upconv2(up2_feat)
  135.         # print(up2_feat.size())
  136.         up3_feat = self.up3(up2_feat, down2_feat)
  137.         # print(up3_feat.size())
  138.         up3_feat = self.upconv3(up3_feat)
  139.         # print(up3_feat.size())
  140.         up4_feat = self.up4(up3_feat, down1_feat)
  141.         # print(up4_feat.size())
  142.         up4_feat = self.upconv4(up4_feat)
  143.         # print(up4_feat.size())
  144.  
  145.         if return_features:
  146.             outputs = up4_feat
  147.         else:
  148.             outputs = self.final(up4_feat)
  149.  
  150.         return outputs
  151.  
  152. class Conv3x3(nn.Module):
  153.     def __init__(self, in_feat, out_feat):
  154.         super(Conv3x3, self).__init__()
  155.  
  156.         self.conv1 = nn.Sequential(nn.Conv2d(in_feat, out_feat,
  157.                                              kernel_size=3,
  158.                                              stride=1,
  159.                                              padding=1),
  160.                                    nn.BatchNorm2d(out_feat),
  161.                                    nn.ReLU())
  162.  
  163.         self.conv2 = nn.Sequential(nn.Conv2d(out_feat, out_feat,
  164.                                              kernel_size=3,
  165.                                              stride=1,
  166.                                              padding=1),
  167.                                    nn.BatchNorm2d(out_feat),
  168.                                    nn.ReLU())
  169.  
  170.     def forward(self, inputs):
  171.         outputs = self.conv1(inputs)
  172.         outputs = self.conv2(outputs)
  173.         return outputs
  174.  
  175.  
  176. class Conv3x3Drop(nn.Module):
  177.     def __init__(self, in_feat, out_feat):
  178.         super(Conv3x3Drop, self).__init__()
  179.  
  180.         self.conv1 = nn.Sequential(nn.Conv2d(in_feat, out_feat,
  181.                                              kernel_size=3,
  182.                                              stride=1,
  183.                                              padding=1),
  184.                                    nn.Dropout(p=0.2),
  185.                                    nn.ReLU())
  186.  
  187.         self.conv2 = nn.Sequential(nn.Conv2d(out_feat, out_feat,
  188.                                              kernel_size=3,
  189.                                              stride=1,
  190.                                              padding=1),
  191.                                    nn.BatchNorm2d(out_feat),
  192.                                    nn.ReLU())
  193.  
  194.     def forward(self, inputs):
  195.         outputs = self.conv1(inputs)
  196.         outputs = self.conv2(outputs)
  197.         return outputs
  198.  
  199.  
  200. class Conv3x3Small(nn.Module):
  201.     def __init__(self, in_feat, out_feat):
  202.         super(Conv3x3Small, self).__init__()
  203.  
  204.         self.conv1 = nn.Sequential(nn.Conv2d(in_feat, out_feat,
  205.                                              kernel_size=3,
  206.                                              stride=1,
  207.                                              padding=1),
  208.                                    nn.ELU(),
  209.                                    nn.Dropout(p=0.2))
  210.  
  211.         self.conv2 = nn.Sequential(nn.Conv2d(out_feat, out_feat,
  212.                                              kernel_size=3,
  213.                                              stride=1,
  214.                                              padding=1),
  215.                                    nn.ELU())
  216.  
  217.     def forward(self, inputs):
  218.         outputs = self.conv1(inputs)
  219.         outputs = self.conv2(outputs)
  220.         return outputs
  221.  
  222.  
  223. class UpConcat(nn.Module):
  224.     def __init__(self, in_feat, out_feat):
  225.         super(UpConcat, self).__init__()
  226.  
  227.         self.up = nn.UpsamplingBilinear2d(scale_factor=2)
  228.  
  229.         # self.deconv = nn.ConvTranspose2d(in_feat, out_feat,
  230.         #                                  kernel_size=3,
  231.         #                                  stride=1,
  232.         #                                  dilation=1)
  233.  
  234.         self.deconv = nn.ConvTranspose2d(in_feat,
  235.                                          out_feat,
  236.                                          kernel_size=2,
  237.                                          stride=2)
  238.  
  239.     def forward(self, inputs, down_outputs):
  240.         # TODO: Upsampling required after deconv?
  241.         # outputs = self.up(inputs)
  242.         outputs = self.deconv(inputs)
  243.         out = torch.cat([down_outputs, outputs], 1)
  244.         return out
  245.  
  246.  
  247. class UpSample(nn.Module):
  248.     def __init__(self, in_feat, out_feat):
  249.         super(UpSample, self).__init__()
  250.  
  251.         self.up = nn.Upsample(scale_factor=2, mode='nearest')
  252.  
  253.         self.deconv = nn.ConvTranspose2d(in_feat,
  254.                                          out_feat,
  255.                                          kernel_size=2,
  256.                                          stride=2)
  257.  
  258.     def forward(self, inputs, down_outputs):
  259.         # TODO: Upsampling required after deconv?
  260.         outputs = self.up(inputs)
  261.         # outputs = self.deconv(inputs)
  262.         out = torch.cat([outputs, down_outputs], 1)
  263.         return out
  264.  
  265. '''
  266. ########################################## STEP 4: INSTANTIATE MODEL CLASS
  267. '''
  268. model = UNet()
  269.  
  270. #####################
  271. # USE GPU FOR MODEL #
  272. #####################
  273.  
  274. if torch.cuda.is_available():
  275.     model.cuda()
  276.  
  277. '''
  278. ########################################## STEP 5: INSTANTIATE LOSS CLASS
  279. '''
  280. criterion = nn.CrossEntropyLoss()
  281.  
  282. '''
  283. ########################################## STEP 6: INSTANTIATE OPTIMIZER CLASS
  284. '''
  285. learning_rate = 0.01 #note: 1 iteration is 100 images, which is our batch size. We update parameters every 100 images
  286.  
  287. optimizer = torch.optim.SGD(model.parameters(), lr = learning_rate)
  288.  
  289. '''
  290. ########################################## STEP 7: TRAIN THE MODEL
  291. '''
  292. iter = 0
  293. for epoch in range(num_epochs):
  294.     for i, (images, labels) in enumerate(train_loader):
  295.         #####################
  296.         # USE GPU FOR MODEL #
  297.         #####################
  298.         if torch.cuda.is_available():
  299.             images = images.unsqueeze(1).type(torch.FloatTensor).cuda()
  300.             labels = labels.cuda()
  301.         else:
  302.             images = images.unsqueeze(1).type(torch.FloatTensor)
  303.        
  304.         # Clear gradients w.r.t. parameters
  305.         optimizer.zero_grad() #don't want to accumulate gradients from previous iterations. Why?
  306.        
  307.         # Forward pass to get ouput/logits
  308.         # Size of outputs is 100 x 10 because each image has output of a value for each digit. Higher value = more likely.
  309.         outputs = model(images)
  310.        
  311.         # Calculate Loss: softmax --> cross entropy loss
  312.         loss = criterion(outputs, labels)
  313.        
  314.         # Getting gradients w.r.t. parameters
  315.         loss.backward()
  316.        
  317.         # Updating parameters
  318.         optimizer.step()
  319.        
  320.         iter += 1
  321.        
  322.         if (iter % 50 == 0):
  323.             # Calculate Accuracy, for every 50 iterations
  324.             correct = 0
  325.             total = 0
  326.             # Iterate through the test dataset
  327.             for images, labels in test_loader:
  328.                 #####################
  329.                 # USE GPU FOR MODEL #
  330.                 #####################
  331.                 if torch.cuda.is_available():
  332.                     images = images.unsqueeze(1).type(torch.FloatTensor).cuda() #deals with dimension issue
  333.  
  334.                 # Forward pass only to get logits/output
  335.                 outputs = model(images)
  336.                
  337.                 # Get predictions from the maximum value
  338.                 _, predicted = torch.max(outputs.data, 1) #need to review how this syntax works
  339.                
  340.                 # Total number of lables
  341.                 total += labels.size(0)
  342.                
  343.                 #####################
  344.                 # USE GPU FOR MODEL #
  345.                 #####################
  346.                 # Total correct predictions... need to bring predicted back to cpu to be able to use .sum() python function
  347.                 if torch.cuda.is_available():
  348.                     correct += (predicted.cpu() == labels.cpu()).sum().item()
  349.                 else:
  350.                     correct += (predicted == labels).sum().item()
  351.                    
  352.             accuracy = 100 * (correct / total)
  353.            
  354.             # Print Loss
  355.             print('Iteration: {}. Loss: {}. Accuracy: {}'.format(iter, loss.item(), accuracy))
  356.  
  357. ''' ERRORS
  358. ---------------------------------------------------------------------------
  359. RuntimeError                              Traceback (most recent call last)
  360. <ipython-input-38-2baa99b3f165> in <module>
  361.     22
  362.     23         # Calculate Loss: softmax --> cross entropy loss
  363. ---> 24         loss = criterion(outputs, labels)
  364.     25
  365.     26         # Getting gradients w.r.t. parameters
  366.  
  367. ~\Anaconda3\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
  368.    487             result = self._slow_forward(*input, **kwargs)
  369.    488         else:
  370. --> 489             result = self.forward(*input, **kwargs)
  371.    490         for hook in self._forward_hooks.values():
  372.    491             hook_result = hook(self, input, result)
  373.  
  374. ~\Anaconda3\lib\site-packages\torch\nn\modules\loss.py in forward(self, input, target)
  375.    902     def forward(self, input, target):
  376.    903         return F.cross_entropy(input, target, weight=self.weight,
  377. --> 904                                ignore_index=self.ignore_index, reduction=self.reduction)
  378.    905
  379.    906
  380.  
  381. ~\Anaconda3\lib\site-packages\torch\nn\functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction)
  382.   1968     if size_average is not None or reduce is not None:
  383.   1969         reduction = _Reduction.legacy_get_string(size_average, reduce)
  384. -> 1970     return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  385.   1971
  386.   1972
  387.  
  388. ~\Anaconda3\lib\site-packages\torch\nn\functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
  389.   1790         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
  390.   1791     elif dim == 4:
  391. -> 1792         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
  392.   1793     else:
  393.   1794         # dim == 3 or dim > 4
  394.  
  395. RuntimeError: Expected object of scalar type Long but got scalar type Byte for argument #2 'target'
  396. '''
Advertisement
RAW Paste Data Copied
Advertisement