实现流程:
一、训练模型
准备数据集,编辑代码进行训练。下载内容在文章末尾。
import torch
import torch.nn as nn
import torchvision.models as models
from torchvision import datasets, transforms
from torch.autograd import Variable
import torch.optim as optim
import os
from trainlog import acc_lossHistory
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
class resnet18(nn.Module):
def __init__(self, num_classes=2):
super(resnet18, self).__init__()
net = models.resnet18(pretrained=True)
# net.classifier = nn.Sequential()
net.fc=nn.Sequential()