python 分类 step优化

# -*- coding: utf-8 -*-
from __future__ import print_function, division
import argparse
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import dataloader
from torchvision import transforms

import time
import os
import platform
import socket
import sys
sys.path.append('../../')

from data.market1501 import Market1501
from data.market_similar import Market_csv_similar
from nets.mobileone import make_mobileone_s0
from nets.model import ft_net, ft_net_dense
from models.skipnet import get_skip
import torch
from torch.nn.functional import one_hot



y_loss = {}  # loss history
y_loss['train'] = []
y_loss['val'] = []
y_err = {}
y_err['train'] = []
y_err['val'] = []

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = model.state_dict()