Torch自查使用

数据集Dataset&Dataloader

dataset = MyDataset（file）

dataloader = DataLoader(dataset, batch_size, shuffle = True)

模型结构 class Model

class MyModel(nn.Module):
    def __init__(self):
        super(MyModule,self).__init__()
        self.net =  nn.Sequential(
            nn.Linear(),
            nn.Sigmoid(),
            nn.Linear()
        )

    def forward(self,x):
        return self.net(x)



class MyModel(nn.Module):
    def __init__(self):
        super(MyModule,self).__init__()
        self.layer1 = nn.Linear(),
        self.layer2 = nn.Sigmoid(),
        self.layer3 = nn.Linear()

    def forward(self,x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        return out

整体流程

//训练
dataset = MyDataset(file)
tr_set = DataLoader(dataset,16,shuffle=True)
model = MyModel.to(device)
criterion = nn.MSEloss()
optimizer = torch.optim.SGD(model.parameter(),0.1)

for epoch in range(n_epochs):
    model.train()
    for x,y in tr.set:
        optimizer.zero_grad()
        x,y = x.to(device), y.to(device)
        pred = model(x)  //forward pass
        loss = criterion(pred,y)
        loss.backward()
        optimizer.step()


//测试
model.eval()
total_loss = 0
for x,y in dv_set:
    x,y = x.to(device), y.to(device)
    with torch.no_grad():
        pred = model(x)
        loss = criterion(pred,y)
    total_loss += loss.cpu().item() * len(x)
    avg_loss = total_loss / len(dv_set.dataset)
    

//验证
model.eval()
preds = []
for x in tt_set：
    x = x.to(device)
    with torch.no_grad():
        pred = model(x)
        preds.append(pred.cpu())


torch.save(model.state_dict(),path)

ckpt = torch.load(path)
model.load_state_dict(ckpt)

加载预训练模型

model=CNN(in_channels=3) #加载模型结构
model.load_state_dict(torch.load('model.pth')) #加载模型参数
model.cuda() #模型加载到GPU
model.eval() #不启用BatchNormalization和Dropout