diff --git a/level1/train.py b/level1/train.py index e2c729e..29565f5 100644 --- a/level1/train.py +++ b/level1/train.py @@ -32,7 +32,7 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25): running_corrects = 0.0 # Iterate over data. - for data in dataloders[phase]: + for data in dataloaders[phase]: # Corrected the spelling from 'dataloders' to 'dataloaders' # get the inputs inputs, labels = data @@ -57,7 +57,7 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25): optimizer.step() # statistics - running_loss += loss.data[0] + running_loss += loss.item() # Changed to .item() for correct loss extraction running_corrects += torch.sum(preds == labels.data).to(torch.float32) epoch_loss = running_loss / dataset_sizes[phase] @@ -82,38 +82,38 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25): if __name__ == '__main__': - # data_transform, pay attention that the input of Normalize() is Tensor and the input of RandomResizedCrop() or RandomHorizontalFlip() is PIL Image + # Data transformations data_transforms = { 'train': transforms.Compose([ - transforms.RandomSizedCrop(224), + transforms.RandomResizedCrop(224), # Changed from RandomSizedCrop to RandomResizedCrop transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ - transforms.Scale(256), + transforms.Resize(256), # Changed from Scale to Resize transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } - # your image data file + # Your image data file data_dir = '/data' image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']} - # wrap your data and label into Tensor - dataloders = {x: torch.utils.data.DataLoader(image_datasets[x], + # Wrap your data and label into Tensor + dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=4) for x in ['train', 'val']} dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} - # use gpu or not + # Use GPU or not use_gpu = torch.cuda.is_available() - # get model and replace the original fc layer with your fc layer + # Get model and replace the original fc layer with your fc layer model_ft = models.resnet18(pretrained=True) num_ftrs = model_ft.fc.in_features model_ft.fc = nn.Linear(num_ftrs, 2) @@ -121,7 +121,7 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25): if use_gpu: model_ft = model_ft.cuda() - # define loss function + # Define loss function criterion = nn.CrossEntropyLoss() # Observe that all parameters are being optimized