Python卷积神经网络resnet训练代码，请给个详细的注释？

Question

这是对人脸图片数据集的训练代码，用的是卷积神经网络resnet，请给以下代码写个详细的注释，最好是行注释。

import torch
import torch.nn as nn
from PIL import Image
from torch import optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class EmoDataset(Dataset):
    def __init__(self, txt_path, transform=None):
        f = open(txt_path, 'r')
        data = f.readlines()
        imgs, labels = [], []
        for line in data:
            word = line.rstrip().split()
            imgs.append(word[0])
            labels.append(int(word[1]))
        self.img = imgs
        self.label = labels
        self.transform = transform

    def __getitem__(self, idx):
        img = self.img[idx]
        label = self.label[idx]
        img = Image.open(img).convert('RGB')
        if transforms is not None:
            img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.label)


def Net(class_num):
    model = models.resnet34(pretrained=False)
    num_ftrs = model.fc.in_features
    # 添加dim=1语句，避免警告
    model.fc = nn.Sequential(
        nn.Dropout(0.5),
        nn.Linear(num_ftrs, class_num),
        nn.Softmax(dim=1)
    )
    model.to(device)
    return model


def Emodataloader(data_txt_path, batch_size):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])])
    dataset = EmoDataset(txt_path=data_txt_path, transform=transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return dataloader





if __name__ == '__main__':
    train()

Answer 1

关注我，注释不难写