{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":6799,"databundleVersionId":4225553,"sourceType":"competition"}],"dockerImageVersionId":30302,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Code for the blog post [Train Resnet50 on ImageNet with PyTorch](https://moiseevigor.github.io/software/2022/12/18/one-pager-training-resnet-on-imagenet/).\n\nThis code will train Resnet50 model on the ImageNet dataset for 10 epochs using ADAM optimizer with a learning rate of 0.001. The model is trained on GPU if available, otherwise it is trained on CPU.","metadata":{}},{"cell_type":"code","source":"import torch\nimport torchvision\nimport torchvision.transforms as transforms\n\n# Set device\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Set hyperparameters\nnum_epochs = 10\nbatch_size = 64\nlearning_rate = 0.001\n\n# Initialize transformations for data augmentation\ntransform = transforms.Compose([\n    transforms.Resize(256),\n    transforms.RandomHorizontalFlip(),\n    transforms.RandomVerticalFlip(),\n    transforms.RandomRotation(degrees=45),\n    transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5),\n    transforms.CenterCrop(224),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])\n\n# Load the ImageNet Object Localization Challenge dataset\ntrain_dataset = '/kaggle/input/imagenet-object-localization-challenge/ILSVRC/Data/CLS-LOC/train'\nprint(f'Loading the ImageNet dataset from {train_dataset}')\ntrain_dataset = torchvision.datasets.ImageFolder(\n    root=train_dataset, \n    transform=transform\n)\n\ntrain_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)\n\n# Load the ResNet50 model\nmodel = torchvision.models.resnet50(pretrained=True)\n\n# Parallelize training across multiple GPUs\nmodel = torch.nn.DataParallel(model)\n\n# Set the model to run on the device\nmodel = model.to(device)\n\n# Define the loss function and optimizer\ncriterion = torch.nn.CrossEntropyLoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n\n# Train the model...\nprint(f'Training the model on ImageNet')\nfor epoch in range(num_epochs):\n    for inputs, labels in train_loader:\n        # Move input and label tensors to the device\n        inputs = inputs.to(device)\n        labels = labels.to(device)\n\n        # Zero out the optimizer\n        optimizer.zero_grad()\n\n        # Forward pass\n        outputs = model(inputs)\n        loss = criterion(outputs, labels)\n\n        # Backward pass\n        loss.backward()\n        optimizer.step()\n\n    # Print the loss for every epoch\n    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}')\n\nprint(f'Finished Training, Loss: {loss.item():.4f}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}