{"cells":[{"metadata":{},"cell_type":"markdown","source":"## 0.fix seeds"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport numpy as np\nimport random\nfrom torch.utils.data.dataset import Dataset\nimport os\nfrom PIL import Image\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\nimport torchvision\nfrom tqdm.notebook import tqdm\nimport time\n\n\ndef set_seed(seed, backends=False):\n    random.seed(seed)\n    np.random.seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\n    # Remove randomness (may be slower on Tesla GPUs)\n    # https://pytorch.org/docs/stable/notes/randomness.html\n    if backends:\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"set_seed(1234)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 1.building dataset"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"class AnimalDataset(Dataset):\n    def __init__(self, root_dir, img_path, transform, class_dict=None, test=False, train=True):\n        self.root_dir = root_dir\n        self.img_path = img_path\n        self.transform = transform\n        self.test = test\n        self.train = train\n        self.class_dict = class_dict\n        self.path_lis = [os.path.join(self.root_dir, self.img_path, pa) for pa in\n                         os.listdir(os.path.join(self.root_dir, img_path))]\n        self.path_lis = sorted(self.path_lis, key=lambda x: int(x.split(\".\")[-2]))\n        self.path_lis = np.random.permutation(self.path_lis)\n        #print(self.path_lis[0])\n        if self.test:\n            self.img_lis = self.path_lis\n        elif self.train:\n            self.img_lis = self.path_lis[:int(len(self.path_lis) * 0.8)]\n        else:\n            self.img_lis = self.path_lis[int(len(self.path_lis) * 0.8):]\n        self.len = len(self.img_lis)\n\n    def __getitem__(self, index):\n        img = Image.open(self.img_lis[index]).convert('RGB')\n        if self.transform is not None:\n            img = self.transform(img)\n        if self.test:\n            return img\n        else:\n            label = self.class_dict[self.img_lis[index].replace(\"\\\\\", \"/\").split(\".\")[-3].split(\"/\")[-1]]\n            return img, label\n\n    def __len__(self):\n        return self.len","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 2.setting transforms"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torchvision.transforms as transforms\n# augmentation\nclass TransForm():\n    def get_train_transforms(self):\n        train_trans = transforms.Compose([\n            transforms.Resize(size=(256, 256)),\n            transforms.RandomCrop(size=(224, 224)),\n            transforms.RandomApply([\n                transforms.RandomHorizontalFlip(p=0.5),\n                transforms.RandomVerticalFlip(p=0.5)\n            ]),\n            transforms.RandomGrayscale(p=0.5),\n            transforms.ColorJitter(brightness=0.3, contrast=0.3, hue=0.1),\n            transforms.ToTensor(),\n            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n        ])\n        return train_trans\n\n    def get_val_transforms(self):\n        val_trans = transforms.Compose([\n            transforms.Resize(size=(224, 224)),\n            transforms.ToTensor(),\n            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n        ])\n        return val_trans\n\n    def get_test_transforms(self):\n        test_trans = transforms.Compose([\n            transforms.Resize(size=(224, 224)),\n            transforms.ToTensor(),\n            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n        ])\n        return test_trans","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 3.train functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def parse_config(config_path):\n    file = open(config_path, mode=\"r\", encoding=\"utf-8\")\n    cfg = yaml.load(file)\n    file.close()\n    return cfg\n\n\ndef train_epoch(model, optimizer, scheduler, device, epoch):\n    print(f\"[log]: [start training epoch {epoch}]\")\n    # train\n    total_size = 0.\n    total_loss = 0.\n    correct = 0.\n    model = model.to(device)\n    model.train()\n    start = time.time()\n    for i,data in tqdm(enumerate(train_loader)):\n        inputs, labels = data\n        inputs = inputs.to(device)\n        labels = labels.to(device)\n        # forward()\n        outputs = model(inputs)\n\n        #backward()\n        optimizer.zero_grad()\n        loss = criterion(outputs, labels.to(device))\n        loss.backward()\n\n        #update weights\n        optimizer.step()\n\n        #统计分类情况\n        _, predicted = torch.max(outputs.data, 1)\n        total_size += labels.size(0)\n        correct += (predicted == labels).squeeze().sum().cpu().numpy()\n        total_loss += loss.item()\n        batch_avg_loss = loss.item() / labels.shape[0]\n\n        if debug:\n            print(f\"[debug]: epoch {epoch} | batch {i} | avg loss : {batch_avg_loss} \")\n\n        # 学习率调整\n    scheduler.step()\n    mid = time.time()\n    print(f\"[log]: epoch:{epoch} | train loss:{round(total_loss/ total_size,5)} | accuracy:{round(correct/ total_size,5)} | train time:{round((mid - start)/60, 2)} min\")\n\n    \n    #val\n    val_size = 0.\n    total_val_loss = 0.\n    val_correct = 0.\n    model.eval()\n    with torch.no_grad():\n        for i, data in enumerate(val_loader):\n            inputs, labels = data\n            inputs = inputs.to(device)\n            labels = labels.to(device)\n            # forward()\n            outputs = model(inputs)\n            loss = criterion(outputs, labels)\n            _, predicted = torch.max(outputs.data, 1)\n            val_size += labels.size(0)\n            val_correct += (predicted == labels).squeeze().sum().cpu().numpy()\n            total_val_loss += loss.item()\n    epoch_avg_loss = round(total_val_loss / val_size, 5)\n    epoch_val_acc = round(val_correct / val_size, 5)\n    end = time.time()\n    print(f\"[log]: epoch:{epoch} | val loss:{epoch_avg_loss} |  accuracy:{epoch_val_acc} | val time:{round((end - mid)/60, 2)} min\")\n    return epoch_avg_loss, epoch_val_acc","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 4.train model"},{"metadata":{"trusted":true},"cell_type":"code","source":"set_seed(1234, backends=False)\nmax_epoch = 20\nbatch_size = 128\nlr = 0.001 \nmomentum = 0.9\nweight_decay = 1e-4\nlabel_dict = {\"cat\": 0, \"dog\": 1}\nroot_dir = \"../input/dogs-vs-cats\"\ncheckpoint_dir = \"./checkpoints\"\ntrain_img_dir = \"train/train\"\ndevice = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\noptimizer = torch.optim.Adam\nscheduler = torch.optim.lr_scheduler.StepLR\ndebug = False\nprint(f\"[log]: setting computation on :{device}\")\n# ============================ step 2/6 数据 ============================\nif not os.path.exists(root_dir):\n    raise Exception(\"ERROR: root dir not exists...\")\nif not os.path.exists(checkpoint_dir):\n    print(\"[warn]: checkpoint dir not exists, creating new directory...\")\n    os.mkdir(\"checkpoints\")\ntrans = TransForm()\ntrainset = AnimalDataset(root_dir, train_img_dir, transform=trans.get_train_transforms(), class_dict=label_dict, train=True)\nvalset = AnimalDataset(root_dir, train_img_dir, transform=trans.get_val_transforms(), class_dict=label_dict, train=False)\nif debug:\n    print(f\"[debug]: train image and label:{trainset[0]}\")\ntrain_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True)\nval_loader = DataLoader(valset, batch_size=batch_size)\n\n# ============================ step 3/6 模型 ============================\n# model = LeNet_bn(classes=2).to(DEVICE)\nmodel = torchvision.models.resnet34(pretrained=True)\nmodel = model.to(device)\n# print(model)\nmodel.fc = nn.Linear(512, out_features=len(label_dict), bias=True)\nprint(f\"[log]: model initialized...\")\n\n# ============================ step 4/6 损失函数 ============================\ncriterion = nn.CrossEntropyLoss().to(device)\n\n# ============================ step 5/6 优化器 +学习策略 ============================\n# optimizer = optim.SGD(model.parameters(), lr=LR, momentum=MOMENTUM)\noptimizer = optimizer(model.parameters(), lr=lr, weight_decay=weight_decay)\nscheduler = scheduler(optimizer, step_size=10, gamma=0.1)\n\n# ============================ step 6/6 训练 ============================\n\nbest_loss = 10000\nfor epoch in range(max_epoch):\n    epoch_avg_loss, epoch_val_acc = train_epoch(model, optimizer, scheduler, device, epoch)\n    # save model\n    torch.save(model.state_dict(), f\"checkpoints/epoch-{epoch}-val-loss-{epoch_avg_loss}.pth\")\n    if epoch_avg_loss < best_loss:\n        torch.save(model.state_dict(), f\"checkpoints/best-model.pth\")\n        best_loss = epoch_avg_loss","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}