Привет. Я пытаюсь обучить модель, которая берет последовательности изображений и классифицирует их. Я написал свою собственную функцию набора данных для загрузки данных, однако у меня есть некоторые запутанные проблемы с памятью. Я прочитал пару папок с последовательностями изображений и добавил их, чтобы сделать пакет. Однако, когда я проверил память с помощью команды nvidia-smi, память графического процессора увеличилась до неожиданного уровня.
Теперь я закомментировал части загрузки изображений, чтобы найти, в чем проблема. Теперь мой набор данных возвращает только один тензор torch.zero с размером [num_frames, channel_num, height, widht]. Я понял, что если num_frames - фиксированное число, память не увеличивается. Но если num_frames отличается в каждой партии памяти GPU начинает увеличиваться. Я понял, что он увеличивается и остается там всякий раз, когда появляется больший тензор num_frames. Даже если следующие партии меньше, память не уменьшается. Так что изменение num_frames создает проблему, как-то. Но поскольку в каждой папке у меня разное количество изображений, мне нужно, чтобы оно менялось. Это команда nvidia-smi? Разве это не показывает реальное использование памяти? Если это правильно, как я могу решить это? Использование памяти фиксированных и различных значений num_frames действительно огромно. Ниже вы можете увидеть мой полный код.
from __future__ import print_function, division
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
import torchvision
from torchvision import datasets, models, transforms,utils
import matplotlib.pyplot as plt
import time
import os
import copy
import torch.nn.functional as F
import pandas as pd
from tqdm import tqdm
from datetime import datetime
import matplotlib.pyplot as plt
from PIL import Image
from skimage import io, transform
from torch.utils.data import Dataset, DataLoader
import h5py
from torch.nn.utils.rnn import pad_sequence
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.image_conv =nn.Sequential(
nn.Conv2d(3, 64, 3, padding = (1,1)),
nn.Conv2d(64, 64, 3, padding = (1,1)),
nn.MaxPool2d(2),
nn.ReLU(),
nn.Conv2d(64, 128, 3, padding = (1,1)),
nn.Conv2d(128, 128, 3, padding = (1,1)),
nn.MaxPool2d(2),
nn.ReLU(),
nn.Conv2d(128, 64, 3, padding = (1,1)),
nn.MaxPool2d(2),
nn.ReLU()
)
self.fc_image = nn.Linear(8*8*64, 4)
def forward(self, x):
#print(x.shape)
x = self.image_conv(x)
x = self.fc_image(x.view((-1,8*8*64)))
return x
def get_label(name): #converts names to int labels
if "ANG" in name:
return 0
elif "HAP" in name:
return 1
elif "NEU" in name:
return 2
elif "SAD" in name:
return 3
count = 0
class Audio_video_dataset(Dataset):
def __init__(self, h5_file, video_root_dir, transform=None):
self.h5_file = h5_file
with h5py.File(self.h5_file, 'r') as f:
self.keys = list(f.keys())
self.video_root_dir = video_root_dir
self.transform = transform
def __len__(self):
return len(self.keys)
def read_images(self, video_dir):
image_list=os.listdir(video_dir)
"""
X = []
for i in range(1,len(image_list)+1):
image = Image.open(os.path.join(video_dir,'image-{:d}.jpeg'.format(i)))
if self.transform:
image = self.transform(image)
X.append(image)
X = torch.stack(X, dim=0)
"""
X =torch.zeros((len(image_list),3,64,64))
print(X.shape)
return X
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
global count
count+=1
selected_elem = self.keys[idx] #current elemen (audio,video)
sequence_name = os.path.join(self.video_root_dir,
selected_elem)
video_dir = sequence_name+".flv"
label = get_label(sequence_name)
image_seq = self.read_images(video_dir)
return image_seq,label
def customBatchBuilder(samples):
image_seq, label = zip(*samples)
image_seq = pad_sequence(image_seq)
image_seq = image_seq.view((-1,3,64,64))
label = torch.Tensor(label).int()
return image_seq,label
dataset_train = Audio_video_dataset(h5_file="audio_features_4class.hdf5",
video_root_dir='cropped_face_frames',
transform=transforms.Compose([
transforms.Resize((64,64)),
transforms.ToTensor()
]))
batch_size = 1
train_loader = DataLoader(dataset_train, batch_size=batch_size,
shuffle=True, num_workers=0, collate_fn=customBatchBuilder)
train_set_size = len(dataset_train)
#valid_set_size = len(dataset_valid)
print("Train set size:",train_set_size)
#print("Test set size:",valid_set_size)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # checks if there is gpu available
print(device)
def train_model(model, criterion, optimizer, num_epochs=25,checkp_epoch=0):
since = time.time()
pbar=tqdm(range(checkp_epoch,num_epochs))
for epoch in pbar: #range(checkp_epoch,num_epochs):
model.train() # Set model to training mode
running_loss = 0.0
running_corrects = 0
for sample in train_loader:
images ,labels = sample
images = images.to(device)
batch_size = images.size(0)
labels =torch.full((batch_size,), labels[0]).long().to(device)
optimizer.zero_grad()
with torch.set_grad_enabled(True):
outputs = model(images)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# statistics
running_loss += loss.item() * images.size(0)
running_corrects += torch.sum(preds == labels.data)
train_loss = running_loss / train_set_size
train_acc = running_corrects.double() / train_set_size
running_loss = 0.0
running_corrects = 0
pbar.set_description("train acc {:.3} loss {:.4}".format(train_acc,train_loss))
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
# load best model weights
return model
model_ft = Net().to(device)
criterion = nn.CrossEntropyLoss()
optimizer_ft = optim.Adam(model_ft.parameters(), lr=1e-4)
num_epochs=100
model_ft = train_model(model_ft, criterion, optimizer_ft,num_epochs=num_epochs)