pytorch实现对输入超过三通道的数据进行训练

 import torch from torch.utils import data import os from PIL import Image import numpy as np import pickle class Fire_Unfire(data.Dataset): def __init__(self,fire_path,unfire_path): self.pickle_fire = open(fire_path,'rb') self.pickle_unfire = open(unfire_path,'rb') def __getitem__(self,index): if index =2504 and index<6489: unfire=pickle.load(self.pickle_unfire) unfire.transpose(2,0,1) data=torch.from_numpy(unfire) label=0 return data,label def __len__(self): 6489 <pre></div><div class="gaodaimacode"><pre class="prettyprint linenums"> root_path = './datasets/train' dataset = Fire_Unfire(root_path +'/fire_train.pkl',root_path +'/unfire_train.pkl') #转换成pytorch网络输入的格式(批量大小,通道数,高,宽) from torch.utils.data import DataLoader fire_dataloader = DataLoader(dataset,batch_size=4,shuffle=True,drop_last = True) 

 import torch from torch.utils import data from nets.mobilenet import mobilenet from config.config import default_config from torch.autograd import Variable as V import numpy as np import sys import time opt = default_config() def train(): #模型定义 model = mobilenet().cuda() if opt.pretrain_model: model.load_state_dict(torch.load(opt.pretrain_model)) #损失函数 criterion = torch.nn.CrossEntropyLoss().cuda() #学习率 lr = opt.lr #优化器 optimizer = torch.optim.SGD(model.parameters(),lr = lr,weight_decay=opt.weight_decay) pre_loss = 0.0 #训练 for epoch in range(opt.max_epoch): #训练数据 train_data = Fire_Unfire(opt.root_path +'/fire_train.pkl',opt.root_path +'/unfire_train.pkl') train_dataloader = data.DataLoader(train_data,batch_size=opt.batch_size,shuffle=True,drop_last = True) loss_sum = 0.0 for i,(datas,labels) in enumerate(train_dataloader): #print(i,datas.size(),labels) #梯度清零 optimizer.zero_grad() #输入 input = V(datas.cuda()).float() #目标 target = V(labels.cuda()).long() #输出 score = model(input).cuda() #损失 loss = criterion(score,target) loss_sum += loss #反向传播 loss.backward() #梯度更新 optimizer.step() print('{}{}{}{}{}'.format('epoch:',epoch,',','loss:',loss)) torch.save(model.state_dict(),'models/mobilenet_%d.pth'%(epoch+370))